@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a69402aaa8

[Perf] Integrate FlashInfer KDA kernels (#55364)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

Wei Zhao committed Sep 7, 2026 at 23:29 UTC a69402aaa8171851232fd6cdd19a33c2a49fd670

8 files changed +722 -155

tests/models/kimi_k3/test_kda.py

+348 -113

@@ -13,6 +13,7 @@ import torch

13 import torch.nn.functional as F

15 from vllm import _custom_ops as ops

16 +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture

17 from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update

18 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (

19 gather_initial_states,

@@ -22,8 +23,11 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (

23 )

24 from vllm.models.kimi_k3.nvidia import kda as nvidia_kda

25 from vllm.models.kimi_k3.nvidia.kda import (

26 + _flashinfer_kda_prefill,

27 _flashkda_prefill,

28 _store_cache_checkpoints_kernel,

29 + is_flashinfer_fused_kda_decode_supported,

30 + is_flashinfer_recurrent_kda_prefill_supported,

31 is_flashkda_supported,

32 is_fused_kda_decode_supported,

33 )

@@ -43,6 +47,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (

47 )

48 from vllm.platforms import current_platform

49 from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd

50 +from vllm.utils.flashinfer import flashinfer_fused_kda_decode

51 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

53 DEVICE = current_platform.device_type

@@ -87,6 +92,7 @@ def test_kda_recoverssm_config_state_layout():

92 ),

93 cache_config=SimpleNamespace(

94 mamba_cache_dtype="auto",

95 + mamba_ssm_cache_dtype="auto",

96 use_kda_recoverssm=True,

97 ),

98 parallel_config=SimpleNamespace(tensor_parallel_size=1),

@@ -909,6 +915,7 @@ def test_kda_recoverssm_verify_and_group_commit(

915 (96, 1, -5.0, True, "DS"),

916 ],

917 )

918 +@pytest.mark.parametrize("decode_backend", ["native", "flashinfer"])

919 @torch.inference_mode()

920 def test_fused_kda_decode_correctness(

921 num_heads: int,

@@ -916,17 +923,36 @@ def test_fused_kda_decode_correctness(

923 lower_bound: float | None,

924 fuse_output_norm: bool,

925 conv_layout: str,

926 + decode_backend: str,

927 ):

928 D, W = 128, 4

921 - if not is_fused_kda_decode_supported(

922 - num_heads,

923 - D,

924 - W,

925 - num_spec=0,

926 - input_dtype=torch.bfloat16,

927 - conv_state_dtype=torch.bfloat16,

928 - ):

929 - pytest.skip("Fused KDA decode is not supported on this platform")

929 + state_dtype = torch.bfloat16 if decode_backend == "flashinfer" else torch.float32

930 + if decode_backend == "flashinfer":

931 + if conv_layout == "DS":

932 + pytest.skip("FlashInfer fused decode requires SD conv-state layout")

933 + if not fuse_output_norm:

934 + pytest.skip("FlashInfer's fused decode always applies output norm")

935 + supported = is_flashinfer_fused_kda_decode_supported(

936 + num_heads,

937 + D,

938 + W,

939 + num_spec=0,

940 + input_dtype=torch.bfloat16,

941 + conv_state_dtype=torch.bfloat16,

942 + recurrent_state_dtype=state_dtype,

943 + )

944 + else:

945 + supported = is_fused_kda_decode_supported(

946 + num_heads,

947 + D,

948 + W,

949 + num_spec=0,

950 + input_dtype=torch.bfloat16,

951 + conv_state_dtype=torch.bfloat16,

952 + recurrent_state_dtype=state_dtype,

953 + )

954 + if not supported:

955 + pytest.skip(f"{decode_backend} fused KDA decode is not supported")

956 torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS"))

957 dim = num_heads * D

958 slots = num_seqs + 2

@@ -995,7 +1021,7 @@ def test_fused_kda_decode_correctness(

1021 D,

1022 dtype=torch.float32,

1023 device=DEVICE,

998 - )

1024 + ).to(state_dtype)

1025

1026 conv_ref = conv_seed.clone()

1027 state_ref = state_seed.clone()

@@ -1030,7 +1056,7 @@ def test_fused_kda_decode_correctness(

1056 conv_slot_elements = 3 * dim * (W - 1)

1057 state_slot_elements = num_heads * D * D

1058 conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize

1033 - page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize

1059 + page_bytes = conv_slot_bytes + state_slot_elements * state_dtype.itemsize

1060 cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE)

1061 conv_actual = torch.as_strided(

1062 cache_storage.view(torch.bfloat16),

@@ -1042,30 +1068,57 @@ def test_fused_kda_decode_correctness(

1068 ),

1069 )

1070 state_actual = torch.as_strided(

1045 - cache_storage.view(torch.float32),

1071 + cache_storage.view(state_dtype),

1072 size=(slots, num_heads, D, D),

1047 - stride=(page_bytes // torch.float32.itemsize, D * D, D, 1),

1048 - storage_offset=conv_slot_bytes // torch.float32.itemsize,

1073 + stride=(page_bytes // state_dtype.itemsize, D * D, D, 1),

1074 + storage_offset=conv_slot_bytes // state_dtype.itemsize,

1075 )

1076 conv_actual.copy_(conv_seed)

1077 state_actual.copy_(state_seed)

1078 fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous()

1053 - actual = ops.fused_kda_decode(

1054 - x=packed_x,

1055 - weight=fused_weight,

1056 - bias=None,

1057 - conv_state=conv_actual,

1058 - raw_g=raw_g,

1059 - raw_beta=raw_beta,

1060 - A_log=A_log,

1061 - dt_bias=dt_bias,

1062 - state_indices=state_indices,

1063 - state=state_actual,

1064 - lower_bound=lower_bound,

1065 - output_gate=output_gate if fuse_output_norm else None,

1066 - norm_weight=norm_weight if fuse_output_norm else None,

1067 - norm_eps=norm_eps,

1068 - )

1079 + if decode_backend == "flashinfer":

1080 + output = torch.empty(

1081 + 1,

1082 + num_seqs,

1083 + num_heads,

1084 + D,

1085 + dtype=torch.bfloat16,

1086 + device=DEVICE,

1087 + )

1088 +

1089 + actual = flashinfer_fused_kda_decode(

1090 + x=packed_x,

1091 + weight=fused_weight,

1092 + conv_state=conv_actual,

1093 + raw_gate=raw_g,

1094 + raw_beta=raw_beta,

1095 + A_log=A_log,

1096 + dt_bias=dt_bias,

1097 + state_indices=state_indices,

1098 + state=state_actual,

1099 + output_gate=output_gate,

1100 + norm_weight=norm_weight,

1101 + lower_bound=lower_bound,

1102 + norm_eps=norm_eps,

1103 + output=output,

1104 + )

1105 + else:

1106 + actual = ops.fused_kda_decode(

1107 + x=packed_x,

1108 + weight=fused_weight,

1109 + bias=None,

1110 + conv_state=conv_actual,

1111 + raw_g=raw_g,

1112 + raw_beta=raw_beta,

1113 + A_log=A_log,

1114 + dt_bias=dt_bias,

1115 + state_indices=state_indices,

1116 + state=state_actual,

1117 + lower_bound=lower_bound,

1118 + output_gate=output_gate if fuse_output_norm else None,

1119 + norm_weight=norm_weight if fuse_output_norm else None,

1120 + norm_eps=norm_eps,

1121 + )

1122

1123 torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)

1124 torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0)

@@ -1080,17 +1133,171 @@ def test_fused_kda_decode_rejects_speculative_conv_state():

1133 num_spec=2,

1134 input_dtype=torch.bfloat16,

1135 conv_state_dtype=torch.bfloat16,

1136 + recurrent_state_dtype=torch.float32,

1137 + )

1138 +

1139 +

1140 +def _make_kda_prefill_inputs(

1141 + state_dtype: torch.dtype,

1142 + *,

1143 + lower_bound: float,

1144 +) -> SimpleNamespace:

1145 + B, T, H, D = 1, 48, 2, 128

1146 + torch.manual_seed(11)

1147 + q, k, v, raw_g = [

1148 + torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)

1149 + ]

1150 + raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)

1151 + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5

1152 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1

1153 + initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE).to(

1154 + state_dtype

1155 + )

1156 + cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)

1157 + return SimpleNamespace(

1158 + q=q,

1159 + k=k,

1160 + v=v,

1161 + raw_g=raw_g,

1162 + raw_beta=raw_beta,

1163 + A_log=A_log,

1164 + dt_bias=dt_bias,

1165 + initial_state=initial_state,

1166 + cu_seqlens=cu_seqlens,

1167 + lower_bound=lower_bound,

1168 + )

1169 +

1170 +

1171 +def _require_kda_prefill_backend(

1172 + backend: str,

1173 + state_dtype: torch.dtype,

1174 + lower_bound: float,

1175 +) -> None:

1176 + if backend == "flashinfer":

1177 + supported = is_flashinfer_recurrent_kda_prefill_supported(

1178 + 128, torch.bfloat16, state_dtype, lower_bound

1179 + )

1180 + else:

1181 + assert backend == "flashkda"

1182 + supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound)

1183 + if not supported:

1184 + pytest.skip(f"{backend} KDA prefill is not supported on this platform")

1185 +

1186 +

1187 +def _run_kda_prefill_backend(

1188 + backend: str,

1189 + *,

1190 + q: torch.Tensor,

1191 + k: torch.Tensor,

1192 + v: torch.Tensor,

1193 + raw_g: torch.Tensor,

1194 + raw_beta: torch.Tensor,

1195 + A_log: torch.Tensor,

1196 + dt_bias: torch.Tensor,

1197 + initial_state: torch.Tensor,

1198 + cu_seqlens: torch.Tensor,

1199 + lower_bound: float,

1200 +) -> tuple[torch.Tensor, torch.Tensor]:

1201 + initial_state = initial_state.clone()

1202 + output = torch.empty_like(v)

1203 + if backend == "flashinfer":

1204 + flashinfer_query_start_loc = cu_seqlens.to(torch.int64)

1205 + seq_order = None

1206 + if q.shape[1] > initial_state.shape[0]:

1207 + seq_order = torch.argsort(

1208 + flashinfer_query_start_loc.diff(), descending=True

1209 + ).to(torch.int32)

1210 + return _flashinfer_kda_prefill(

1211 + q=q,

1212 + k=k,

1213 + v=v,

1214 + raw_g=raw_g,

1215 + raw_beta=raw_beta,

1216 + A_log=A_log,

1217 + dt_bias=dt_bias,

1218 + lower_bound=lower_bound,

1219 + initial_state=initial_state,

1220 + cu_seqlens=flashinfer_query_start_loc,

1221 + out=output,

1222 + seq_order=seq_order,

1223 + )

1224 +

1225 + assert backend == "flashkda"

1226 + import vllm._flashkda_C # noqa: F401

1227 +

1228 + final_state = torch.empty_like(initial_state)

1229 + workspace = torch.empty(

1230 + torch.ops._flashkda_C.get_workspace_size(

1231 + q.shape[1], q.shape[2], cu_seqlens.numel() - 1

1232 + ),

1233 + dtype=torch.uint8,

1234 + device=q.device,

1235 + )

1236 + return _flashkda_prefill(

1237 + q=q,

1238 + k=k,

1239 + v=v,

1240 + g=raw_g,

1241 + beta=raw_beta,

1242 + A_log=A_log,

1243 + dt_bias=dt_bias,

1244 + lower_bound=lower_bound,

1245 + initial_state=initial_state,

1246 + cu_seqlens=cu_seqlens,

1247 + out=output,

1248 + final_state=final_state,

1249 + workspace=workspace,

1250 + )

1251 +

1252 +

1253 +def _kda_prefill_reference(

1254 + inputs: SimpleNamespace,

1255 +) -> tuple[torch.Tensor, torch.Tensor]:

1256 + gate = inputs.lower_bound * torch.sigmoid(

1257 + inputs.A_log.exp()[None, None, :, None]

1258 + * (inputs.raw_g.float() + inputs.dt_bias[None, None])

1259 + )

1260 + beta = inputs.raw_beta.float().sigmoid()

1261 + q_norm = l2norm_fwd(inputs.q.contiguous())

1262 + k_norm = l2norm_fwd(inputs.k.contiguous())

1263 + expected_outputs = []

1264 + expected_states = []

1265 + for i, (start, end) in enumerate(

1266 + zip(inputs.cu_seqlens[:-1].tolist(), inputs.cu_seqlens[1:].tolist())

1267 + ):

1268 + output, final_state = naive_recurrent_kda(

1269 + q_norm[:, start:end],

1270 + k_norm[:, start:end],

1271 + inputs.v[:, start:end],

1272 + gate[:, start:end],

1273 + beta[:, start:end],

1274 + initial_state=inputs.initial_state[i : i + 1].transpose(-1, -2),

1275 + output_final_state=True,

1276 + )

1277 + assert final_state is not None

1278 + expected_outputs.append(output)

1279 + expected_states.append(final_state)

1280 + return (

1281 + torch.cat(expected_outputs, dim=1),

1282 + torch.cat(expected_states).transpose(-1, -2).contiguous(),

1283 )

1284

1285

1286 +@pytest.mark.parametrize(

1287 + ("backend", "state_dtype"),

1288 + [

1289 + pytest.param("flashkda", torch.float32, id="flashkda"),

1290 + pytest.param("flashinfer", torch.bfloat16, id="flashinfer"),

1291 + ],

1292 +)

1293 @torch.inference_mode()

1087 -def test_flashkda_near_collinear_keys_remain_finite():

1294 +def test_kda_prefill_near_collinear_keys_remain_finite(

1295 + backend: str,

1296 + state_dtype: torch.dtype,

1297 +):

1298 """Guard against unstable inversion of near-collinear key blocks."""

1299 lower_bound = -5.0

1090 - if not is_flashkda_supported(128, torch.bfloat16, lower_bound):

1091 - pytest.skip("FlashKDA is not supported on this platform")

1092 -

1093 - import vllm._flashkda_C # noqa: F401

1300 + _require_kda_prefill_backend(backend, state_dtype, lower_bound)

1301

1302 T, H, D = 16384, 1, 128

1303 torch.manual_seed(0)

@@ -1102,81 +1309,131 @@ def test_flashkda_near_collinear_keys_remain_finite():

1309 raw_beta = torch.full((1, T, H), 8.0, dtype=qk.dtype, device=DEVICE)

1310 A_log = torch.zeros(H, dtype=torch.float32, device=DEVICE)

1311 dt_bias = torch.zeros(H, D, dtype=torch.float32, device=DEVICE)

1105 - initial_state = torch.zeros(1, H, D, D, dtype=torch.float32, device=DEVICE)

1106 - final_state = torch.empty_like(initial_state)

1107 - output = torch.empty_like(value)

1312 + initial_state = torch.zeros(1, H, D, D, dtype=state_dtype, device=DEVICE)

1313 cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=DEVICE)

1109 - workspace = torch.empty(

1110 - torch.ops._flashkda_C.get_workspace_size(T, H, 1),

1111 - dtype=torch.uint8,

1112 - device=DEVICE,

1113 - )

1114 -

1115 - torch.ops._flashkda_C.fwd(

1116 - qk,

1117 - qk,

1118 - value,

1119 - raw_gate,

1120 - raw_beta,

1121 - D**-0.5,

1122 - output,

1123 - workspace,

1124 - A_log,

1125 - dt_bias,

1126 - lower_bound,

1127 - initial_state,

1128 - final_state,

1129 - cu_seqlens,

1314 + output, final_state = _run_kda_prefill_backend(

1315 + backend,

1316 + q=qk,

1317 + k=qk,

1318 + v=value,

1319 + raw_g=raw_gate,

1320 + raw_beta=raw_beta,

1321 + A_log=A_log,

1322 + dt_bias=dt_bias,

1323 + initial_state=initial_state,

1324 + cu_seqlens=cu_seqlens,

1325 + lower_bound=lower_bound,

1326 )

1327

1328 assert torch.isfinite(output).all()

1329 assert torch.isfinite(final_state).all()

1330

1331

1332 +@pytest.mark.parametrize(

1333 + ("backend", "state_dtype", "tolerance"),

1334 + [

1335 + pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"),

1336 + pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"),

1337 + pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"),

1338 + ],

1339 +)

1340 @torch.inference_mode()

1137 -def test_flashkda_correctness():

1138 - if not is_flashkda_supported(128, torch.bfloat16, -3.0):

1139 - pytest.skip("FlashKDA is not supported on this platform")

1341 +def test_kda_prefill_correctness(

1342 + backend: str,

1343 + state_dtype: torch.dtype,

1344 + tolerance: float,

1345 +):

1346 + lower_bound = -5.0

1347 + _require_kda_prefill_backend(backend, state_dtype, lower_bound)

1348 + inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound)

1349 + expected_out, expected_state = _kda_prefill_reference(inputs)

1350 + actual_out, actual_state = _run_kda_prefill_backend(

1351 + backend,

1352 + **vars(inputs),

1353 + )

1354

1141 - import vllm._flashkda_C # noqa: F401

1355 + assert_close("o", expected_out, actual_out, tolerance)

1356 + assert_close("ht", expected_state, actual_state, tolerance)

1357

1143 - B, T, H, D = 1, 48, 2, 128

1144 - torch.manual_seed(11)

1358 +

1359 +@torch.inference_mode()

1360 +def test_flashinfer_kda_prefill_breakable_graph_cross_stream():

1361 + if not is_flashinfer_recurrent_kda_prefill_supported(

1362 + 128,

1363 + torch.bfloat16,

1364 + torch.bfloat16,

1365 + -5.0,

1366 + ):

1367 + pytest.skip("FlashInfer KDA prefill is not supported on this platform")

1368 +

1369 + B, T, H, D = 1, 8, 12, 128

1370 q, k, v, raw_g = [

1371 torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)

1372 ]

1148 - beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)

1149 - A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5

1150 - dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1

1151 - initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE)

1152 - cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)

1373 + raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)

1374 + initial_state = torch.randn(

1375 + B,

1376 + H,

1377 + D,

1378 + D,

1379 + dtype=torch.bfloat16,

1380 + device=DEVICE,

1381 + )

1382 + kwargs = {

1383 + "q": q,

1384 + "k": k,

1385 + "v": v,

1386 + "raw_g": raw_g,

1387 + "raw_beta": raw_beta,

1388 + "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE),

1389 + "dt_bias": torch.randn(H, D, dtype=torch.float32, device=DEVICE),

1390 + "lower_bound": -5.0,

1391 + "initial_state": initial_state,

1392 + "cu_seqlens": torch.tensor([0, T], dtype=torch.int64, device=DEVICE),

1393 + "out": torch.empty_like(q),

1394 + "seq_order": torch.zeros(B, dtype=torch.int32, device=DEVICE),

1395 + }

1396 +

1397 + capture_stream = torch.Stream(device=DEVICE)

1398 + original_stream = torch.accelerator.current_stream()

1399 + torch.accelerator.set_stream(capture_stream)

1400 + try:

1401 + graph_value = torch.zeros(1, device=DEVICE)

1402 + capture = BreakableCUDAGraphCapture()

1403 + with capture:

1404 + graph_value.add_(1)

1405 + capture.add_eager(lambda: _flashinfer_kda_prefill(**kwargs))

1406 + graph_value.add_(1)

1407 + capture_stream.synchronize()

1408 + finally:

1409 + torch.accelerator.set_stream(original_stream)

1410 +

1411 + torch.testing.assert_close(graph_value, torch.zeros_like(graph_value))

1412 + capture.replay()

1413 + torch.accelerator.synchronize()

1414 + torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2))

1415 +

1416 +

1417 +@torch.inference_mode()

1418 +def test_flashkda_checkpoint_correctness():

1419 lower_bound = -3.0

1420 + _require_kda_prefill_backend("flashkda", torch.float32, lower_bound)

1421 +

1422 + import vllm._flashkda_C # noqa: F401

1423

1424 + inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound)

1425 + q, k, v = inputs.q, inputs.k, inputs.v

1426 + raw_g, raw_beta = inputs.raw_g, inputs.raw_beta

1427 + A_log, dt_bias = inputs.A_log, inputs.dt_bias

1428 + initial_state, cu_seqlens = inputs.initial_state, inputs.cu_seqlens

1429 + _, T, H, D = q.shape

1430 + expected_out, expected_state = _kda_prefill_reference(inputs)

1431 gate = lower_bound * torch.sigmoid(

1432 A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :])

1433 )

1158 - beta = beta_logits.float().sigmoid()

1434 + beta = raw_beta.float().sigmoid()

1435 q_norm = l2norm_fwd(q.contiguous())

1436 k_norm = l2norm_fwd(k.contiguous())

1161 -

1162 - expected_outputs = []

1163 - expected_states = []

1164 - for i, (start, end) in enumerate(

1165 - zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist())

1166 - ):

1167 - output, final_state = naive_recurrent_kda(

1168 - q_norm[:, start:end],

1169 - k_norm[:, start:end],

1170 - v[:, start:end],

1171 - gate[:, start:end],

1172 - beta[:, start:end],

1173 - initial_state=initial_state[i].transpose(-1, -2),

1174 - output_final_state=True,

1175 - )

1176 - expected_outputs.append(output)

1177 - expected_states.append(final_state)

1178 - expected_out = torch.cat(expected_outputs, dim=1)

1179 - expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()

1437 _, expected_checkpoint = naive_recurrent_kda(

1438 q_norm[:, :16],

1439 k_norm[:, :16],

@@ -1189,33 +1446,11 @@ def test_flashkda_correctness():

1446 assert expected_checkpoint is not None

1447 expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous()

1448

1192 - actual_out = torch.empty_like(v)

1193 - actual_state = torch.empty_like(initial_state)

1449 workspace = torch.empty(

1450 torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1),

1451 dtype=torch.uint8,

1452 device=DEVICE,

1453 )

1199 - torch.ops._flashkda_C.fwd(

1200 - q,

1201 - k,

1202 - v,

1203 - raw_g,

1204 - beta_logits,

1205 - D**-0.5,

1206 - actual_out,

1207 - workspace,

1208 - A_log,

1209 - dt_bias,

1210 - lower_bound,

1211 - initial_state,

1212 - actual_state,

1213 - cu_seqlens,

1214 - )

1215 -

1216 - assert_close("o", expected_out, actual_out, 0.01)

1217 - assert_close("ht", expected_state, actual_state, 0.01)

1218 -

1454 checkpoint_out = torch.empty_like(v)

1455 checkpoint_final_state = torch.empty_like(initial_state)

1456 checkpoint_state = torch.empty_like(initial_state)

@@ -1225,7 +1460,7 @@ def test_flashkda_correctness():

1460 k=k,

1461 v=v,

1462 g=raw_g,

1228 - beta=beta_logits,

1463 + beta=raw_beta,

1464 A_log=A_log,

1465 dt_bias=dt_bias,

1466 lower_bound=lower_bound,

tests/models/kimi_k3/test_kda_metadata.py

+4 -1

@@ -49,6 +49,8 @@ PRUNED_METADATA_FIELDS = {

49 "prefill_state_indices",

50 "prefill_has_initial_state",

51 "spec_sequence_masks",

52 + "flashinfer_prefill_query_start_loc",

53 + "flashinfer_prefill_seq_order",

54 }

@@ -59,10 +61,10 @@ def _assert_matches_shared_gdn(

61 assert actual.recoverssm_context is None

62 for field in fields(GDNAttentionMetadata):

63 actual_value = getattr(actual, field.name)

62 - expected_value = getattr(reference, field.name)

64 if field.name in PRUNED_METADATA_FIELDS:

65 assert actual_value is None

66 continue

67 + expected_value = getattr(reference, field.name)

68 if (

69 field.name in {"spec_token_indx", "non_spec_token_indx"}

70 and actual.num_spec_decodes > 0

@@ -157,6 +159,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):

159 ),

160 cache_config=SimpleNamespace(

161 mamba_cache_dtype="auto",

162 + mamba_ssm_cache_dtype="auto",

163 use_kda_recoverssm=True,

164 ),

165 parallel_config=SimpleNamespace(tensor_parallel_size=1),

vllm/engine/arg_utils.py

+14 -2

@@ -772,7 +772,10 @@ class EngineArgs:

772

773 fail_on_environ_validation: bool = False

774 gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None

775 - kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None

775 + kda_prefill_backend: Literal["auto", "triton", "flashkda", "flashinfer"] | None = (

776 + None

777 + )

778 + kda_decode_backend: Literal["auto", "native", "flashinfer", "triton"] | None = None

779

780 def __post_init__(self):

781 # support `EngineArgs(compilation_config={...})`

@@ -1747,10 +1750,17 @@ class EngineArgs:

1750 parser.add_argument(

1751 "--kda-prefill-backend",

1752 dest="kda_prefill_backend",

1750 - choices=["auto", "triton", "flashkda"],

1753 + choices=["auto", "triton", "flashkda", "flashinfer"],

1754 default=None,

1755 help="Select KDA prefill backend.",

1756 )

1757 + parser.add_argument(

1758 + "--kda-decode-backend",

1759 + dest="kda_decode_backend",

1760 + choices=["auto", "native", "flashinfer", "triton"],

1761 + default=None,

1762 + help="Select KDA decode backend.",

1763 + )

1764 return parser

1765

1766 @classmethod

@@ -2567,6 +2577,8 @@ class EngineArgs:

2577 self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend

2578 if self.kda_prefill_backend is not None:

2579 self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend

2580 + if self.kda_decode_backend is not None:

2581 + self.additional_config["kda_decode_backend"] = self.kda_decode_backend

2582

2583 config = VllmConfig(

2584 model_config=model_config,

vllm/model_executor/layers/mamba/mamba_utils.py

+12 -2

@@ -134,9 +134,19 @@ class MambaStateDtypeCalculator:

134 cls,

135 model_dtype: ModelDType | torch.dtype,

136 mamba_cache_dtype: MambaDType,

137 + mamba_ssm_cache_dtype: MambaDType = "auto",

138 ) -> tuple[torch.dtype, torch.dtype]:

138 - state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)

139 - return (state_dtype, torch.float32)

139 + conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)

140 + if mamba_ssm_cache_dtype == "auto":

141 + recurrent_state_dtype = torch.float32

142 + elif mamba_ssm_cache_dtype in ("float32", "bfloat16"):

143 + recurrent_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype]

144 + else:

145 + raise ValueError(

146 + "KDA recurrent state supports only auto, float32, and bfloat16; "

147 + f"got {mamba_ssm_cache_dtype}"

148 + )

149 + return (conv_state_dtype, recurrent_state_dtype)

150

151 @classmethod

152 def append_kda_recoverssm_record(

vllm/models/kimi_k3/nvidia/kda.py

+284 -36

@@ -50,6 +50,12 @@ from vllm.platforms import current_platform

50 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated

51 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

52 from vllm.triton_utils import tl, triton

53 +from vllm.utils.flashinfer import (

54 + flashinfer_fused_kda_decode,

55 + flashinfer_recurrent_kda,

56 + has_flashinfer_fused_kda_decode,

57 + has_flashinfer_recurrent_kda,

58 +)

59 from vllm.v1.attention.backend import AttentionBackend

60 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

61 from vllm.v1.kv_cache_interface import MambaSpec

@@ -155,6 +161,7 @@ def is_fused_kda_decode_supported(

161 num_spec: int,

162 input_dtype: torch.dtype,

163 conv_state_dtype: torch.dtype,

164 + recurrent_state_dtype: torch.dtype,

165 ) -> bool:

166 # The fused kernel handles both conv-state cache layouts (SD and DS); the

167 # inner strides are selected from the tensor at launch time.

@@ -165,6 +172,7 @@ def is_fused_kda_decode_supported(

172 or num_spec != 0

173 or input_dtype != torch.bfloat16

174 or conv_state_dtype != torch.bfloat16

175 + or recurrent_state_dtype != torch.float32

176 or not hasattr(torch.ops._C, "fused_kda_decode")

177 ):

178 return False

@@ -176,9 +184,93 @@ def is_fused_kda_decode_supported(

184 )

185

186

187 +def is_flashinfer_fused_kda_decode_supported(

188 + num_heads: int,

189 + head_dim: int,

190 + conv_width: int,

191 + num_spec: int,

192 + input_dtype: torch.dtype,

193 + conv_state_dtype: torch.dtype,

194 + recurrent_state_dtype: torch.dtype,

195 +) -> bool:

196 + if not has_flashinfer_fused_kda_decode():

197 + return False

198 + capability = current_platform.get_device_capability()

199 + if capability is None:

200 + return False

201 + compute_capability = (capability.major, capability.minor)

202 + return (

203 + compute_capability in ((10, 0), (10, 3))

204 + and num_heads in (12, 24, 32, 48, 96)

205 + and head_dim == 128

206 + and conv_width == 4

207 + and num_spec == 0

208 + and input_dtype == torch.bfloat16

209 + and conv_state_dtype == torch.bfloat16

210 + and recurrent_state_dtype in (torch.float32, torch.bfloat16)

211 + and not is_conv_state_dim_first()

212 + )

213 +

214 +

215 +def resolve_kda_decode_backend(

216 + backend: str,

217 + num_heads: int,

218 + head_dim: int,

219 + conv_width: int,

220 + num_spec: int,

221 + input_dtype: torch.dtype,

222 + conv_state_dtype: torch.dtype,

223 + recurrent_state_dtype: torch.dtype,

224 +) -> str:

225 + if backend not in ("auto", "native", "flashinfer", "triton"):

226 + raise ValueError(f"Unsupported KDA decode backend: {backend}")

227 + if backend == "triton":

228 + return "triton"

229 +

230 + native_supported = is_fused_kda_decode_supported(

231 + num_heads,

232 + head_dim,

233 + conv_width,

234 + num_spec,

235 + input_dtype,

236 + conv_state_dtype,

237 + recurrent_state_dtype,

238 + )

239 + if native_supported and backend in ("auto", "native"):

240 + logger.info_once("Using native fused KDA decode backend.")

241 + return "native"

242 + if backend == "native":

243 + raise RuntimeError(

244 + "Native fused KDA decode requires a supported CUDA architecture, "

245 + "bfloat16 activations and convolution state, float32 recurrent "

246 + "state, head_dim=128, convolution width 4, and no speculation."

247 + )

248 +

249 + flashinfer_supported = is_flashinfer_fused_kda_decode_supported(

250 + num_heads,

251 + head_dim,

252 + conv_width,

253 + num_spec,

254 + input_dtype,

255 + conv_state_dtype,

256 + recurrent_state_dtype,

257 + )

258 + if flashinfer_supported and backend == "flashinfer":

259 + logger.info_once("Using FlashInfer fused KDA decode backend.")

260 + return "flashinfer"

261 + if backend == "flashinfer":

262 + raise RuntimeError(

263 + "FlashInfer fused KDA decode requires CUDA SM100 or SM103, "

264 + "bfloat16 activations and convolution state, bfloat16 or float32 "

265 + "recurrent state, head_dim=128, convolution width 4, no speculation."

266 + )

267 + return "triton"

268 +

269 +

270 def is_flashkda_supported(

271 head_dim: int,

181 - dtype: torch.dtype,

272 + input_dtype: torch.dtype,

273 + recurrent_state_dtype: torch.dtype,

274 lower_bound: float | None,

275 ) -> bool:

276 if not current_platform.is_cuda():

@@ -188,7 +280,29 @@ def is_flashkda_supported(

280 capability is not None

281 and capability.major in (9, 10, 12)

282 and head_dim == 128

191 - and dtype == torch.bfloat16

283 + and input_dtype == torch.bfloat16

284 + and recurrent_state_dtype in (torch.bfloat16, torch.float32)

285 + and lower_bound is not None

286 + )

287 +

288 +

289 +def is_flashinfer_recurrent_kda_prefill_supported(

290 + head_dim: int,

291 + input_dtype: torch.dtype,

292 + recurrent_state_dtype: torch.dtype,

293 + lower_bound: float | None,

294 +) -> bool:

295 + if not current_platform.is_cuda() or not has_flashinfer_recurrent_kda():

296 + return False

297 + capability = current_platform.get_device_capability()

298 + if capability is None:

299 + return False

300 + compute_capability = (capability.major, capability.minor)

301 + return (

302 + compute_capability in ((10, 0), (10, 3))

303 + and head_dim == 128

304 + and input_dtype == torch.bfloat16

305 + and recurrent_state_dtype == torch.bfloat16

306 and lower_bound is not None

307 )

308

@@ -237,6 +351,44 @@ def _flashkda_prefill(

351 return out, final_state

352

353

354 +def _flashinfer_kda_prefill(

355 + q: torch.Tensor,

356 + k: torch.Tensor,

357 + v: torch.Tensor,

358 + raw_g: torch.Tensor,

359 + raw_beta: torch.Tensor,

360 + A_log: torch.Tensor,

361 + dt_bias: torch.Tensor,

362 + lower_bound: float,

363 + initial_state: torch.Tensor,

364 + cu_seqlens: torch.Tensor,

365 + out: torch.Tensor,

366 + seq_order: torch.Tensor | None = None,

367 + prefill_workspace: object | None = None,

368 +) -> tuple[torch.Tensor, torch.Tensor]:

369 + output, _ = flashinfer_recurrent_kda(

370 + q=q.contiguous(),

371 + k=k.contiguous(),

372 + v=v.contiguous(),

373 + g=raw_g.contiguous(),

374 + beta=raw_beta.contiguous(),

375 + A_log=A_log.contiguous(),

376 + dt_bias=dt_bias.contiguous(),

377 + scale=q.shape[-1] ** -0.5,

378 + initial_state=initial_state.contiguous(),

379 + output_final_state=False,

380 + use_qk_l2norm_in_kernel=True,

381 + use_gate_in_kernel=True,

382 + lower_bound=lower_bound,

383 + cu_seqlens=cu_seqlens.contiguous(),

384 + output=out,

385 + beta_is_logit=True,

386 + seq_order=seq_order,

387 + prefill_workspace=prefill_workspace,

388 + )

389 + return output, initial_state

390 +

391 +

392 @triton.jit

393 def _store_cache_checkpoints_kernel(

394 x_ptr,

@@ -303,18 +455,40 @@ def _store_cache_checkpoints_kernel(

455 def resolve_kda_prefill_backend(

456 backend: str,

457 head_dim: int,

306 - dtype: torch.dtype,

458 + input_dtype: torch.dtype,

459 + recurrent_state_dtype: torch.dtype,

460 lower_bound: float | None,

461 ) -> str:

309 - if backend not in ("auto", "triton", "flashkda"):

462 + if backend not in ("auto", "triton", "flashkda", "flashinfer"):

463 raise ValueError(f"Unsupported KDA prefill backend: {backend}")

311 - supported = is_flashkda_supported(head_dim, dtype, lower_bound)

312 - if backend == "flashkda" and not supported:

464 + flashinfer_supported = is_flashinfer_recurrent_kda_prefill_supported(

465 + head_dim,

466 + input_dtype,

467 + recurrent_state_dtype,

468 + lower_bound,

469 + )

470 + flashkda_supported = is_flashkda_supported(

471 + head_dim,

472 + input_dtype,

473 + recurrent_state_dtype,

474 + lower_bound,

475 + )

476 + if backend == "flashinfer" and not flashinfer_supported:

477 raise RuntimeError(

314 - "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, "

315 - "head_dim=128, and a bounded KDA gate."

478 + "FlashInfer KDA prefill requires CUDA SM100 or SM103, bfloat16 "

479 + "activations and recurrent state, head_dim=128, a bounded KDA "

480 + "gate, and flashinfer-python 0.6.18 or newer."

481 )

317 - if supported and backend != "triton":

482 + if backend == "flashkda" and not flashkda_supported:

483 + raise RuntimeError(

484 + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16 input, "

485 + "bfloat16 or float32 recurrent state, head_dim=128, and a bounded "

486 + "KDA gate."

487 + )

488 + if flashinfer_supported and backend == "flashinfer":

489 + logger.info_once("Using FlashInfer KDA prefill backend.")

490 + return "flashinfer"

491 + if flashkda_supported and backend in ("auto", "flashkda"):

492 logger.info_once("Using FlashKDA KDA prefill backend.")

493 return "flashkda"

494 return "triton"

@@ -369,7 +543,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

543 if self.model_config is None or self.cache_config is None:

544 raise ValueError("model_config and cache_config must be set")

545 base_dtypes = MambaStateDtypeCalculator.kda_state_dtype(

372 - self.model_config.dtype, self.cache_config.mamba_cache_dtype

546 + self.model_config.dtype,

547 + self.cache_config.mamba_cache_dtype,

548 + self.cache_config.mamba_ssm_cache_dtype,

549 )

550 if self.cache_config.use_kda_recoverssm:

551 return MambaStateDtypeCalculator.append_kda_recoverssm_record(

@@ -486,17 +662,25 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

662 self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

663 # Keep a width-major copy for fused decode without changing the layout

664 # consumed by the prefill and fallback decode kernels.

489 - conv_state_dtype = self.get_state_dtype()[0]

490 - decode_conv1d_weight = None

491 - if is_fused_kda_decode_supported(

665 + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2]

666 + additional_config = vllm_config.additional_config

667 + decode_backend = (

668 + additional_config.get("kda_decode_backend", "auto")

669 + if isinstance(additional_config, dict)

670 + else "auto"

671 + )

672 + self.kda_decode_backend = resolve_kda_decode_backend(

673 + decode_backend,

674 self.local_num_heads,

675 self.head_dim,

676 self.conv_size,

677 self.num_spec,

678 vllm_config.model_config.dtype,

679 conv_state_dtype,

498 - ):

499 - logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.")

680 + recurrent_state_dtype,

681 + )

682 + decode_conv1d_weight = None

683 + if self.kda_decode_backend != "triton":

684 decode_conv1d_weight = torch.empty(

685 3,

686 self.conv_size,

@@ -533,7 +717,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

717 f"Got {self.gate_lower_bound}."

718 )

719

536 - additional_config = vllm_config.additional_config

720 backend = (

721 additional_config.get("kda_prefill_backend", "auto")

722 if isinstance(additional_config, dict)

@@ -543,11 +726,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

726 backend,

727 self.head_dim,

728 vllm_config.model_config.dtype,

729 + recurrent_state_dtype,

730 self.gate_lower_bound,

731 )

732 self._flashkda_buffer_specs: (

733 tuple[tuple[tuple[int, ...], torch.dtype], ...] | None

734 ) = None

735 + self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = (

736 + None

737 + )

738 if self.kda_prefill_backend == "flashkda":

739 T = vllm_config.scheduler_config.max_num_batched_tokens

740 N = vllm_config.scheduler_config.max_num_seqs

@@ -561,6 +748,13 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

748 ((N, H, D, D), self.get_state_dtype()[1]),

749 ((workspace_size,), torch.uint8),

750 )

751 + elif self.kda_prefill_backend == "flashinfer":

752 + T = vllm_config.scheduler_config.max_num_batched_tokens

753 + H, D = self.local_num_heads, self.head_dim

754 + self._flashinfer_kda_output_spec = (

755 + (1, T, H, D),

756 + self.model_config.dtype,

757 + )

758

759 self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")

760 decode_norm_weight = None

@@ -725,30 +919,50 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

919 conv_state = conv_state.transpose(-1, -2)

920

921 if (

728 - self.decode_conv1d_weight is not None

922 + self.kda_decode_backend != "triton"

923 + and self.decode_conv1d_weight is not None

924 and self.decode_norm_weight is not None

925 and not has_spec_decode

926 and m.num_prefills == 0

927 and m.num_decodes > 0

928 ):

929 assert non_spec_state_indices_tensor is not None

735 - ops.fused_kda_decode(

736 - x=mixed_qkv,

737 - weight=self.decode_conv1d_weight,

738 - bias=self.conv1d.bias,

739 - conv_state=conv_state,

740 - raw_g=g1,

741 - raw_beta=beta,

742 - A_log=self.A_log,

743 - dt_bias=self.dt_bias,

744 - state_indices=non_spec_state_indices_tensor[:num_actual_tokens],

745 - state=recurrent_state,

746 - out=core_attn_out[:, :num_actual_tokens],

747 - lower_bound=self.gate_lower_bound,

748 - output_gate=g2[:num_actual_tokens],

749 - norm_weight=self.decode_norm_weight,

750 - norm_eps=self.o_norm.eps,

751 - )

930 + state_indices = non_spec_state_indices_tensor[:num_actual_tokens]

931 + if self.kda_decode_backend == "flashinfer":

932 + flashinfer_fused_kda_decode(

933 + x=mixed_qkv,

934 + weight=self.decode_conv1d_weight,

935 + conv_state=conv_state,

936 + raw_gate=g1,

937 + raw_beta=beta,

938 + A_log=self.A_log,

939 + dt_bias=self.dt_bias,

940 + state_indices=state_indices,

941 + state=recurrent_state,

942 + output_gate=g2[:num_actual_tokens],

943 + norm_weight=self.decode_norm_weight,

944 + lower_bound=self.gate_lower_bound,

945 + norm_eps=self.o_norm.eps,

946 + output=core_attn_out[:, :num_actual_tokens],

947 + )

948 + else:

949 + ops.fused_kda_decode(

950 + x=mixed_qkv,

951 + weight=self.decode_conv1d_weight,

952 + bias=self.conv1d.bias,

953 + conv_state=conv_state,

954 + raw_g=g1,

955 + raw_beta=beta,

956 + A_log=self.A_log,

957 + dt_bias=self.dt_bias,

958 + state_indices=state_indices,

959 + state=recurrent_state,

960 + out=core_attn_out[:, :num_actual_tokens],

961 + lower_bound=self.gate_lower_bound,

962 + output_gate=g2[:num_actual_tokens],

963 + norm_weight=self.decode_norm_weight,

964 + norm_eps=self.o_norm.eps,

965 + )

966 return

967

968 conv_weights = self.conv1d.weight.view(

@@ -991,6 +1205,35 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1205 final_state=final_state[: initial_state.shape[0]],

1206 workspace=workspace,

1207 )

1208 + elif self.kda_prefill_backend == "flashinfer":

1209 + assert self.gate_lower_bound is not None

1210 + assert m.flashinfer_prefill_query_start_loc is not None

1211 + if q_ns.shape[1] > initial_state.shape[0]:

1212 + assert m.flashinfer_prefill_seq_order is not None

1213 + flashinfer_out = core_attn_out[:, : q_ns.shape[1]]

1214 + if has_spec_decode:

1215 + assert self._flashinfer_kda_output_spec is not None

1216 + (workspace_out,) = current_workspace_manager().get_simultaneous(

1217 + self._flashinfer_kda_output_spec

1218 + )

1219 + flashinfer_out = workspace_out[:, : q_ns.shape[1]]

1220 + (

1221 + core_attn_out_non_spec,

1222 + last_recurrent_state,

1223 + ) = _flashinfer_kda_prefill(

1224 + q=q_ns,

1225 + k=k_ns,

1226 + v=v_ns,

1227 + raw_g=g1_ns,

1228 + raw_beta=beta_ns,

1229 + A_log=self.A_log,

1230 + dt_bias=self.dt_bias,

1231 + lower_bound=self.gate_lower_bound,

1232 + initial_state=initial_state,

1233 + cu_seqlens=m.flashinfer_prefill_query_start_loc,

1234 + out=flashinfer_out,

1235 + seq_order=m.flashinfer_prefill_seq_order,

1236 + )

1237 else:

1238 (

1239 core_attn_out_non_spec,

@@ -1009,7 +1252,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1252 use_qk_l2norm_in_kernel=True,

1253 cu_seqlens=non_spec_query_start_loc,

1254 )

1012 - recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state

1255 + recurrent_state[non_spec_state_indices_tensor] = (

1256 + last_recurrent_state.to(recurrent_state.dtype)

1257 + )

1258 else:

1259 # Pure non-speculative decode.

1260 assert non_spec_state_indices_tensor is not None

@@ -1046,7 +1291,10 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1291 core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

1292 core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)

1293 elif core_attn_out_non_spec is not None:

1049 - if self.kda_prefill_backend != "flashkda" or m.num_prefills == 0:

1294 + if (

1295 + self.kda_prefill_backend not in ("flashkda", "flashinfer")

1296 + or m.num_prefills == 0

1297 + ):

1298 # TODO: decode kernels write directly to core_attn_out

1299 core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[

1300 0, :num_actual_tokens

vllm/models/kimi_k3/nvidia/kda_metadata.py

+23

@@ -270,6 +270,8 @@ class KDACheckpointMetadata:

270

271 @dataclass

272 class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):

273 + flashinfer_prefill_query_start_loc: torch.Tensor | None = None

274 + flashinfer_prefill_seq_order: torch.Tensor | None = None

275 recoverssm_commit: KDARecoverSSMCommitMetadata | None = None

276 recoverssm_context: "KDARecoverSSMCommitContext | None" = field(

277 default=None, repr=False, compare=False

@@ -318,6 +320,11 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

320 device: torch.device,

321 ) -> None:

322 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

323 + additional_config = vllm_config.additional_config

324 + self.use_flashinfer_prefill = (

325 + isinstance(additional_config, dict)

326 + and additional_config.get("kda_prefill_backend") == "flashinfer"

327 + )

328 self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm

329 self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1

330 self.recoverssm_num_accepted_tokens: torch.Tensor | None = None

@@ -721,6 +728,20 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

728 align=align,

729 )

730

731 + flashinfer_prefill_query_start_loc = None

732 + flashinfer_prefill_seq_order = None

733 + if self.use_flashinfer_prefill and num_prefills > 0:

734 + assert non_spec_query_start_loc is not None

735 + flashinfer_prefill_query_start_loc = non_spec_query_start_loc.to(

736 + torch.int64

737 + )

738 + num_non_spec_requests = non_spec_query_start_loc.shape[0] - 1

739 + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens

740 + if num_non_spec_tokens > num_non_spec_requests:

741 + flashinfer_prefill_seq_order = torch.argsort(

742 + flashinfer_prefill_query_start_loc.diff(), descending=True

743 + ).to(torch.int32)

744 +

745 return KimiK3KDAMetadata(

746 num_prefills=num_prefills,

747 num_prefill_tokens=num_prefill_tokens,

@@ -738,6 +759,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

759 spec_token_indx=spec_token_indx,

760 non_spec_token_indx=non_spec_token_indx,

761 num_accepted_tokens=num_accepted_tokens,

762 + flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc,

763 + flashinfer_prefill_seq_order=flashinfer_prefill_seq_order,

764 recoverssm_commit=recoverssm_commit,

765 recoverssm_context=(

766 self._get_recoverssm_context()

vllm/models/kimi_k3/nvidia/model.py

+3 -1

@@ -1649,7 +1649,9 @@ class KimiLinearForCausalLM(

1649 vllm_config: "VllmConfig",

1650 ) -> tuple[torch.dtype, ...]:

1651 dtypes = MambaStateDtypeCalculator.kda_state_dtype(

1652 - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype

1652 + vllm_config.model_config.dtype,

1653 + vllm_config.cache_config.mamba_cache_dtype,

1654 + vllm_config.cache_config.mamba_ssm_cache_dtype,

1655 )

1656 if vllm_config.cache_config.use_kda_recoverssm:

1657 dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record(

vllm/utils/flashinfer.py

+34

@@ -264,6 +264,14 @@ flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper(

264 "flashinfer.decode",

265 "xqa_batch_decode_with_kv_cache",

266 )

267 +flashinfer_recurrent_kda = _lazy_import_wrapper(

268 + "flashinfer.kda",

269 + "recurrent_kda",

270 +)

271 +flashinfer_fused_kda_decode = _lazy_import_wrapper(

272 + "flashinfer.kda_decode",

273 + "fused_kda_decode",

274 +)

275

276

277 # Special case for autotune since it returns a context manager

@@ -399,6 +407,28 @@ def has_flashinfer_bf16_fp4() -> bool:

407 )

408

409

410 +@functools.cache

411 +def has_flashinfer_recurrent_kda() -> bool:

412 + """Return whether FlashInfer recurrent KDA prefill is available."""

413 + if not has_flashinfer():

414 + return False

415 + mod = _get_submodule("flashinfer.kda")

416 + return mod is not None and callable(getattr(mod, "recurrent_kda", None))

417 +

418 +

419 +@functools.cache

420 +def has_flashinfer_fused_kda_decode() -> bool:

421 + """Return whether FlashInfer fused KDA decode is available."""

422 + if not has_flashinfer():

423 + return False

424 + mod = _get_submodule("flashinfer.kda_decode")

425 + return (

426 + mod is not None

427 + and bool(getattr(mod, "_FUSED_KDA_DECODE_AVAILABLE", False))

428 + and callable(getattr(mod, "fused_kda_decode", None))

429 + )

430 +

431 +

432 @functools.cache

433 def has_flashinfer_trtllm_fused_moe() -> bool:

434 """Return `True` if FlashInfer TRTLLM fused MoE is available."""

@@ -1261,6 +1291,8 @@ __all__ = [

1291 "flashinfer_trtllm_batch_decode_with_kv_cache_mla",

1292 "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",

1293 "flashinfer_xqa_batch_decode_with_kv_cache",

1294 + "flashinfer_recurrent_kda",

1295 + "flashinfer_fused_kda_decode",

1296 "autotune",

1297 "has_flashinfer_moe",

1298 "has_flashinfer_comm",

@@ -1268,6 +1300,8 @@ __all__ = [

1300 "has_flashinfer_nvlink_one_sided",

1301 "has_flashinfer_cutlass_fused_moe",

1302 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",

1303 + "has_flashinfer_recurrent_kda",

1304 + "has_flashinfer_fused_kda_decode",

1305 "has_flashinfer_cutedsl_moe_nvfp4",

1306 "has_flashinfer_bf16_fp4",

1307 "has_flashinfer_b12x_moe",