From a69402aaa8171851232fd6cdd19a33c2a49fd670 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Mon, 7 Sep 2026 23:29:30 -0400 Subject: [PATCH] [Perf] Integrate FlashInfer KDA kernels (#55364) Signed-off-by: wzhao18 --- tests/models/kimi_k3/test_kda.py | 463 +++++++++++++----- tests/models/kimi_k3/test_kda_metadata.py | 5 +- vllm/engine/arg_utils.py | 16 +- .../layers/mamba/mamba_utils.py | 14 +- vllm/models/kimi_k3/nvidia/kda.py | 320 ++++++++++-- vllm/models/kimi_k3/nvidia/kda_metadata.py | 23 + vllm/models/kimi_k3/nvidia/model.py | 4 +- vllm/utils/flashinfer.py | 34 ++ 8 files changed, 723 insertions(+), 156 deletions(-) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index e40d2e716d..16d7a9742f 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -13,6 +13,7 @@ import torch import torch.nn.functional as F from vllm import _custom_ops as ops +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, @@ -22,8 +23,11 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import ( ) from vllm.models.kimi_k3.nvidia import kda as nvidia_kda from vllm.models.kimi_k3.nvidia.kda import ( + _flashinfer_kda_prefill, _flashkda_prefill, _store_cache_checkpoints_kernel, + is_flashinfer_fused_kda_decode_supported, + is_flashinfer_recurrent_kda_prefill_supported, is_flashkda_supported, is_fused_kda_decode_supported, ) @@ -43,6 +47,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.utils.flashinfer import flashinfer_fused_kda_decode from vllm.v1.attention.backends.utils import NULL_BLOCK_ID DEVICE = current_platform.device_type @@ -87,6 +92,7 @@ def test_kda_recoverssm_config_state_layout(): ), cache_config=SimpleNamespace( mamba_cache_dtype="auto", + mamba_ssm_cache_dtype="auto", use_kda_recoverssm=True, ), parallel_config=SimpleNamespace(tensor_parallel_size=1), @@ -909,6 +915,7 @@ def test_kda_recoverssm_verify_and_group_commit( (96, 1, -5.0, True, "DS"), ], ) +@pytest.mark.parametrize("decode_backend", ["native", "flashinfer"]) @torch.inference_mode() def test_fused_kda_decode_correctness( num_heads: int, @@ -916,17 +923,36 @@ def test_fused_kda_decode_correctness( lower_bound: float | None, fuse_output_norm: bool, conv_layout: str, + decode_backend: str, ): D, W = 128, 4 - if not is_fused_kda_decode_supported( - num_heads, - D, - W, - num_spec=0, - input_dtype=torch.bfloat16, - conv_state_dtype=torch.bfloat16, - ): - pytest.skip("Fused KDA decode is not supported on this platform") + state_dtype = torch.bfloat16 if decode_backend == "flashinfer" else torch.float32 + if decode_backend == "flashinfer": + if conv_layout == "DS": + pytest.skip("FlashInfer fused decode requires SD conv-state layout") + if not fuse_output_norm: + pytest.skip("FlashInfer's fused decode always applies output norm") + supported = is_flashinfer_fused_kda_decode_supported( + num_heads, + D, + W, + num_spec=0, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + recurrent_state_dtype=state_dtype, + ) + else: + supported = is_fused_kda_decode_supported( + num_heads, + D, + W, + num_spec=0, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + recurrent_state_dtype=state_dtype, + ) + if not supported: + pytest.skip(f"{decode_backend} fused KDA decode is not supported") torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS")) dim = num_heads * D slots = num_seqs + 2 @@ -995,7 +1021,7 @@ def test_fused_kda_decode_correctness( D, dtype=torch.float32, device=DEVICE, - ) + ).to(state_dtype) conv_ref = conv_seed.clone() state_ref = state_seed.clone() @@ -1030,7 +1056,7 @@ def test_fused_kda_decode_correctness( conv_slot_elements = 3 * dim * (W - 1) state_slot_elements = num_heads * D * D conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize - page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize + page_bytes = conv_slot_bytes + state_slot_elements * state_dtype.itemsize cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE) conv_actual = torch.as_strided( cache_storage.view(torch.bfloat16), @@ -1042,30 +1068,57 @@ def test_fused_kda_decode_correctness( ), ) state_actual = torch.as_strided( - cache_storage.view(torch.float32), + cache_storage.view(state_dtype), size=(slots, num_heads, D, D), - stride=(page_bytes // torch.float32.itemsize, D * D, D, 1), - storage_offset=conv_slot_bytes // torch.float32.itemsize, + stride=(page_bytes // state_dtype.itemsize, D * D, D, 1), + storage_offset=conv_slot_bytes // state_dtype.itemsize, ) conv_actual.copy_(conv_seed) state_actual.copy_(state_seed) fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous() - actual = ops.fused_kda_decode( - x=packed_x, - weight=fused_weight, - bias=None, - conv_state=conv_actual, - raw_g=raw_g, - raw_beta=raw_beta, - A_log=A_log, - dt_bias=dt_bias, - state_indices=state_indices, - state=state_actual, - lower_bound=lower_bound, - output_gate=output_gate if fuse_output_norm else None, - norm_weight=norm_weight if fuse_output_norm else None, - norm_eps=norm_eps, - ) + if decode_backend == "flashinfer": + output = torch.empty( + 1, + num_seqs, + num_heads, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + + actual = flashinfer_fused_kda_decode( + x=packed_x, + weight=fused_weight, + conv_state=conv_actual, + raw_gate=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + state=state_actual, + output_gate=output_gate, + norm_weight=norm_weight, + lower_bound=lower_bound, + norm_eps=norm_eps, + output=output, + ) + else: + actual = ops.fused_kda_decode( + x=packed_x, + weight=fused_weight, + bias=None, + conv_state=conv_actual, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + state=state_actual, + lower_bound=lower_bound, + output_gate=output_gate if fuse_output_norm else None, + norm_weight=norm_weight if fuse_output_norm else None, + norm_eps=norm_eps, + ) torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) 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(): num_spec=2, input_dtype=torch.bfloat16, conv_state_dtype=torch.bfloat16, + recurrent_state_dtype=torch.float32, ) +def _make_kda_prefill_inputs( + state_dtype: torch.dtype, + *, + lower_bound: float, +) -> SimpleNamespace: + B, T, H, D = 1, 48, 2, 128 + torch.manual_seed(11) + q, k, v, raw_g = [ + torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) + ] + raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE).to( + state_dtype + ) + cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) + return SimpleNamespace( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + + +def _require_kda_prefill_backend( + backend: str, + state_dtype: torch.dtype, + lower_bound: float, +) -> None: + if backend == "flashinfer": + supported = is_flashinfer_recurrent_kda_prefill_supported( + 128, torch.bfloat16, state_dtype, lower_bound + ) + else: + assert backend == "flashkda" + supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound) + if not supported: + pytest.skip(f"{backend} KDA prefill is not supported on this platform") + + +def _run_kda_prefill_backend( + backend: str, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + lower_bound: float, +) -> tuple[torch.Tensor, torch.Tensor]: + initial_state = initial_state.clone() + output = torch.empty_like(v) + if backend == "flashinfer": + flashinfer_query_start_loc = cu_seqlens.to(torch.int64) + seq_order = None + if q.shape[1] > initial_state.shape[0]: + seq_order = torch.argsort( + flashinfer_query_start_loc.diff(), descending=True + ).to(torch.int32) + return _flashinfer_kda_prefill( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + cu_seqlens=flashinfer_query_start_loc, + out=output, + seq_order=seq_order, + ) + + assert backend == "flashkda" + import vllm._flashkda_C # noqa: F401 + + final_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size( + q.shape[1], q.shape[2], cu_seqlens.numel() - 1 + ), + dtype=torch.uint8, + device=q.device, + ) + return _flashkda_prefill( + q=q, + k=k, + v=v, + g=raw_g, + beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + out=output, + final_state=final_state, + workspace=workspace, + ) + + +def _kda_prefill_reference( + inputs: SimpleNamespace, +) -> tuple[torch.Tensor, torch.Tensor]: + gate = inputs.lower_bound * torch.sigmoid( + inputs.A_log.exp()[None, None, :, None] + * (inputs.raw_g.float() + inputs.dt_bias[None, None]) + ) + beta = inputs.raw_beta.float().sigmoid() + q_norm = l2norm_fwd(inputs.q.contiguous()) + k_norm = l2norm_fwd(inputs.k.contiguous()) + expected_outputs = [] + expected_states = [] + for i, (start, end) in enumerate( + zip(inputs.cu_seqlens[:-1].tolist(), inputs.cu_seqlens[1:].tolist()) + ): + output, final_state = naive_recurrent_kda( + q_norm[:, start:end], + k_norm[:, start:end], + inputs.v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=inputs.initial_state[i : i + 1].transpose(-1, -2), + output_final_state=True, + ) + assert final_state is not None + expected_outputs.append(output) + expected_states.append(final_state) + return ( + torch.cat(expected_outputs, dim=1), + torch.cat(expected_states).transpose(-1, -2).contiguous(), + ) + + +@pytest.mark.parametrize( + ("backend", "state_dtype"), + [ + pytest.param("flashkda", torch.float32, id="flashkda"), + pytest.param("flashinfer", torch.bfloat16, id="flashinfer"), + ], +) @torch.inference_mode() -def test_flashkda_near_collinear_keys_remain_finite(): +def test_kda_prefill_near_collinear_keys_remain_finite( + backend: str, + state_dtype: torch.dtype, +): """Guard against unstable inversion of near-collinear key blocks.""" lower_bound = -5.0 - if not is_flashkda_supported(128, torch.bfloat16, lower_bound): - pytest.skip("FlashKDA is not supported on this platform") - - import vllm._flashkda_C # noqa: F401 + _require_kda_prefill_backend(backend, state_dtype, lower_bound) T, H, D = 16384, 1, 128 torch.manual_seed(0) @@ -1102,81 +1309,131 @@ def test_flashkda_near_collinear_keys_remain_finite(): raw_beta = torch.full((1, T, H), 8.0, dtype=qk.dtype, device=DEVICE) A_log = torch.zeros(H, dtype=torch.float32, device=DEVICE) dt_bias = torch.zeros(H, D, dtype=torch.float32, device=DEVICE) - initial_state = torch.zeros(1, H, D, D, dtype=torch.float32, device=DEVICE) - final_state = torch.empty_like(initial_state) - output = torch.empty_like(value) + initial_state = torch.zeros(1, H, D, D, dtype=state_dtype, device=DEVICE) cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=DEVICE) - workspace = torch.empty( - torch.ops._flashkda_C.get_workspace_size(T, H, 1), - dtype=torch.uint8, - device=DEVICE, - ) - - torch.ops._flashkda_C.fwd( - qk, - qk, - value, - raw_gate, - raw_beta, - D**-0.5, - output, - workspace, - A_log, - dt_bias, - lower_bound, - initial_state, - final_state, - cu_seqlens, + output, final_state = _run_kda_prefill_backend( + backend, + q=qk, + k=qk, + v=value, + raw_g=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, ) assert torch.isfinite(output).all() assert torch.isfinite(final_state).all() +@pytest.mark.parametrize( + ("backend", "state_dtype", "tolerance"), + [ + pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"), + pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"), + pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"), + ], +) @torch.inference_mode() -def test_flashkda_correctness(): - if not is_flashkda_supported(128, torch.bfloat16, -3.0): - pytest.skip("FlashKDA is not supported on this platform") +def test_kda_prefill_correctness( + backend: str, + state_dtype: torch.dtype, + tolerance: float, +): + lower_bound = -5.0 + _require_kda_prefill_backend(backend, state_dtype, lower_bound) + inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound) + expected_out, expected_state = _kda_prefill_reference(inputs) + actual_out, actual_state = _run_kda_prefill_backend( + backend, + **vars(inputs), + ) - import vllm._flashkda_C # noqa: F401 + assert_close("o", expected_out, actual_out, tolerance) + assert_close("ht", expected_state, actual_state, tolerance) - B, T, H, D = 1, 48, 2, 128 - torch.manual_seed(11) + +@torch.inference_mode() +def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): + if not is_flashinfer_recurrent_kda_prefill_supported( + 128, + torch.bfloat16, + torch.bfloat16, + -5.0, + ): + pytest.skip("FlashInfer KDA prefill is not supported on this platform") + + B, T, H, D = 1, 8, 12, 128 q, k, v, raw_g = [ torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) ] - beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) - A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 - dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 - initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE) - cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) - lower_bound = -3.0 + raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) + initial_state = torch.randn( + B, + H, + D, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + kwargs = { + "q": q, + "k": k, + "v": v, + "raw_g": raw_g, + "raw_beta": raw_beta, + "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE), + "dt_bias": torch.randn(H, D, dtype=torch.float32, device=DEVICE), + "lower_bound": -5.0, + "initial_state": initial_state, + "cu_seqlens": torch.tensor([0, T], dtype=torch.int64, device=DEVICE), + "out": torch.empty_like(q), + "seq_order": torch.zeros(B, dtype=torch.int32, device=DEVICE), + } + capture_stream = torch.Stream(device=DEVICE) + original_stream = torch.accelerator.current_stream() + torch.accelerator.set_stream(capture_stream) + try: + graph_value = torch.zeros(1, device=DEVICE) + capture = BreakableCUDAGraphCapture() + with capture: + graph_value.add_(1) + capture.add_eager(lambda: _flashinfer_kda_prefill(**kwargs)) + graph_value.add_(1) + capture_stream.synchronize() + finally: + torch.accelerator.set_stream(original_stream) + + torch.testing.assert_close(graph_value, torch.zeros_like(graph_value)) + capture.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2)) + + +@torch.inference_mode() +def test_flashkda_checkpoint_correctness(): + lower_bound = -3.0 + _require_kda_prefill_backend("flashkda", torch.float32, lower_bound) + + import vllm._flashkda_C # noqa: F401 + + inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound) + q, k, v = inputs.q, inputs.k, inputs.v + raw_g, raw_beta = inputs.raw_g, inputs.raw_beta + A_log, dt_bias = inputs.A_log, inputs.dt_bias + initial_state, cu_seqlens = inputs.initial_state, inputs.cu_seqlens + _, T, H, D = q.shape + expected_out, expected_state = _kda_prefill_reference(inputs) gate = lower_bound * torch.sigmoid( A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :]) ) - beta = beta_logits.float().sigmoid() + beta = raw_beta.float().sigmoid() q_norm = l2norm_fwd(q.contiguous()) k_norm = l2norm_fwd(k.contiguous()) - - expected_outputs = [] - expected_states = [] - for i, (start, end) in enumerate( - zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()) - ): - output, final_state = naive_recurrent_kda( - q_norm[:, start:end], - k_norm[:, start:end], - v[:, start:end], - gate[:, start:end], - beta[:, start:end], - initial_state=initial_state[i].transpose(-1, -2), - output_final_state=True, - ) - expected_outputs.append(output) - expected_states.append(final_state) - expected_out = torch.cat(expected_outputs, dim=1) - expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous() _, expected_checkpoint = naive_recurrent_kda( q_norm[:, :16], k_norm[:, :16], @@ -1189,33 +1446,11 @@ def test_flashkda_correctness(): assert expected_checkpoint is not None expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous() - actual_out = torch.empty_like(v) - actual_state = torch.empty_like(initial_state) workspace = torch.empty( torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1), dtype=torch.uint8, device=DEVICE, ) - torch.ops._flashkda_C.fwd( - q, - k, - v, - raw_g, - beta_logits, - D**-0.5, - actual_out, - workspace, - A_log, - dt_bias, - lower_bound, - initial_state, - actual_state, - cu_seqlens, - ) - - assert_close("o", expected_out, actual_out, 0.01) - assert_close("ht", expected_state, actual_state, 0.01) - checkpoint_out = torch.empty_like(v) checkpoint_final_state = torch.empty_like(initial_state) checkpoint_state = torch.empty_like(initial_state) @@ -1225,7 +1460,7 @@ def test_flashkda_correctness(): k=k, v=v, g=raw_g, - beta=beta_logits, + beta=raw_beta, A_log=A_log, dt_bias=dt_bias, lower_bound=lower_bound, diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 25ba2d6b17..f344b27656 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -49,6 +49,8 @@ PRUNED_METADATA_FIELDS = { "prefill_state_indices", "prefill_has_initial_state", "spec_sequence_masks", + "flashinfer_prefill_query_start_loc", + "flashinfer_prefill_seq_order", } @@ -59,10 +61,10 @@ def _assert_matches_shared_gdn( assert actual.recoverssm_context is None for field in fields(GDNAttentionMetadata): actual_value = getattr(actual, field.name) - expected_value = getattr(reference, field.name) if field.name in PRUNED_METADATA_FIELDS: assert actual_value is None continue + expected_value = getattr(reference, field.name) if ( field.name in {"spec_token_indx", "non_spec_token_indx"} and actual.num_spec_decodes > 0 @@ -157,6 +159,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): ), cache_config=SimpleNamespace( mamba_cache_dtype="auto", + mamba_ssm_cache_dtype="auto", use_kda_recoverssm=True, ), parallel_config=SimpleNamespace(tensor_parallel_size=1), diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 644d7e3139..1a2354a430 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -772,7 +772,10 @@ class EngineArgs: fail_on_environ_validation: bool = False gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None - kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None + kda_prefill_backend: Literal["auto", "triton", "flashkda", "flashinfer"] | None = ( + None + ) + kda_decode_backend: Literal["auto", "native", "flashinfer", "triton"] | None = None def __post_init__(self): # support `EngineArgs(compilation_config={...})` @@ -1747,10 +1750,17 @@ class EngineArgs: parser.add_argument( "--kda-prefill-backend", dest="kda_prefill_backend", - choices=["auto", "triton", "flashkda"], + choices=["auto", "triton", "flashkda", "flashinfer"], default=None, help="Select KDA prefill backend.", ) + parser.add_argument( + "--kda-decode-backend", + dest="kda_decode_backend", + choices=["auto", "native", "flashinfer", "triton"], + default=None, + help="Select KDA decode backend.", + ) return parser @classmethod @@ -2567,6 +2577,8 @@ class EngineArgs: self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend if self.kda_prefill_backend is not None: self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend + if self.kda_decode_backend is not None: + self.additional_config["kda_decode_backend"] = self.kda_decode_backend config = VllmConfig( model_config=model_config, diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 73b1d934e8..8ca1d83955 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -134,9 +134,19 @@ class MambaStateDtypeCalculator: cls, model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, + mamba_ssm_cache_dtype: MambaDType = "auto", ) -> tuple[torch.dtype, torch.dtype]: - state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) - return (state_dtype, torch.float32) + conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) + if mamba_ssm_cache_dtype == "auto": + recurrent_state_dtype = torch.float32 + elif mamba_ssm_cache_dtype in ("float32", "bfloat16"): + recurrent_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype] + else: + raise ValueError( + "KDA recurrent state supports only auto, float32, and bfloat16; " + f"got {mamba_ssm_cache_dtype}" + ) + return (conv_state_dtype, recurrent_state_dtype) @classmethod def append_kda_recoverssm_record( diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 6e33be668e..c5c556d745 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -50,6 +50,12 @@ from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.triton_utils import tl, triton +from vllm.utils.flashinfer import ( + flashinfer_fused_kda_decode, + flashinfer_recurrent_kda, + has_flashinfer_fused_kda_decode, + has_flashinfer_recurrent_kda, +) from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.kv_cache_interface import MambaSpec @@ -155,6 +161,7 @@ def is_fused_kda_decode_supported( num_spec: int, input_dtype: torch.dtype, conv_state_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, ) -> bool: # The fused kernel handles both conv-state cache layouts (SD and DS); the # inner strides are selected from the tensor at launch time. @@ -165,6 +172,7 @@ def is_fused_kda_decode_supported( or num_spec != 0 or input_dtype != torch.bfloat16 or conv_state_dtype != torch.bfloat16 + or recurrent_state_dtype != torch.float32 or not hasattr(torch.ops._C, "fused_kda_decode") ): return False @@ -176,9 +184,93 @@ def is_fused_kda_decode_supported( ) +def is_flashinfer_fused_kda_decode_supported( + num_heads: int, + head_dim: int, + conv_width: int, + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, +) -> bool: + if not has_flashinfer_fused_kda_decode(): + return False + capability = current_platform.get_device_capability() + if capability is None: + return False + compute_capability = (capability.major, capability.minor) + return ( + compute_capability in ((10, 0), (10, 3)) + and num_heads in (12, 24, 32, 48, 96) + and head_dim == 128 + and conv_width == 4 + and num_spec == 0 + and input_dtype == torch.bfloat16 + and conv_state_dtype == torch.bfloat16 + and recurrent_state_dtype in (torch.float32, torch.bfloat16) + and not is_conv_state_dim_first() + ) + + +def resolve_kda_decode_backend( + backend: str, + num_heads: int, + head_dim: int, + conv_width: int, + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, +) -> str: + if backend not in ("auto", "native", "flashinfer", "triton"): + raise ValueError(f"Unsupported KDA decode backend: {backend}") + if backend == "triton": + return "triton" + + native_supported = is_fused_kda_decode_supported( + num_heads, + head_dim, + conv_width, + num_spec, + input_dtype, + conv_state_dtype, + recurrent_state_dtype, + ) + if native_supported and backend in ("auto", "native"): + logger.info_once("Using native fused KDA decode backend.") + return "native" + if backend == "native": + raise RuntimeError( + "Native fused KDA decode requires a supported CUDA architecture, " + "bfloat16 activations and convolution state, float32 recurrent " + "state, head_dim=128, convolution width 4, and no speculation." + ) + + flashinfer_supported = is_flashinfer_fused_kda_decode_supported( + num_heads, + head_dim, + conv_width, + num_spec, + input_dtype, + conv_state_dtype, + recurrent_state_dtype, + ) + if flashinfer_supported and backend == "flashinfer": + logger.info_once("Using FlashInfer fused KDA decode backend.") + return "flashinfer" + if backend == "flashinfer": + raise RuntimeError( + "FlashInfer fused KDA decode requires CUDA SM100 or SM103, " + "bfloat16 activations and convolution state, bfloat16 or float32 " + "recurrent state, head_dim=128, convolution width 4, no speculation." + ) + return "triton" + + def is_flashkda_supported( head_dim: int, - dtype: torch.dtype, + input_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, lower_bound: float | None, ) -> bool: if not current_platform.is_cuda(): @@ -188,7 +280,29 @@ def is_flashkda_supported( capability is not None and capability.major in (9, 10, 12) and head_dim == 128 - and dtype == torch.bfloat16 + and input_dtype == torch.bfloat16 + and recurrent_state_dtype in (torch.bfloat16, torch.float32) + and lower_bound is not None + ) + + +def is_flashinfer_recurrent_kda_prefill_supported( + head_dim: int, + input_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + if not current_platform.is_cuda() or not has_flashinfer_recurrent_kda(): + return False + capability = current_platform.get_device_capability() + if capability is None: + return False + compute_capability = (capability.major, capability.minor) + return ( + compute_capability in ((10, 0), (10, 3)) + and head_dim == 128 + and input_dtype == torch.bfloat16 + and recurrent_state_dtype == torch.bfloat16 and lower_bound is not None ) @@ -237,6 +351,44 @@ def _flashkda_prefill( return out, final_state +def _flashinfer_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor, + seq_order: torch.Tensor | None = None, + prefill_workspace: object | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + output, _ = flashinfer_recurrent_kda( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=raw_g.contiguous(), + beta=raw_beta.contiguous(), + A_log=A_log.contiguous(), + dt_bias=dt_bias.contiguous(), + scale=q.shape[-1] ** -0.5, + initial_state=initial_state.contiguous(), + output_final_state=False, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens.contiguous(), + output=out, + beta_is_logit=True, + seq_order=seq_order, + prefill_workspace=prefill_workspace, + ) + return output, initial_state + + @triton.jit def _store_cache_checkpoints_kernel( x_ptr, @@ -303,18 +455,40 @@ def _store_cache_checkpoints_kernel( def resolve_kda_prefill_backend( backend: str, head_dim: int, - dtype: torch.dtype, + input_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, lower_bound: float | None, ) -> str: - if backend not in ("auto", "triton", "flashkda"): + if backend not in ("auto", "triton", "flashkda", "flashinfer"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") - supported = is_flashkda_supported(head_dim, dtype, lower_bound) - if backend == "flashkda" and not supported: + flashinfer_supported = is_flashinfer_recurrent_kda_prefill_supported( + head_dim, + input_dtype, + recurrent_state_dtype, + lower_bound, + ) + flashkda_supported = is_flashkda_supported( + head_dim, + input_dtype, + recurrent_state_dtype, + lower_bound, + ) + if backend == "flashinfer" and not flashinfer_supported: raise RuntimeError( - "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, " - "head_dim=128, and a bounded KDA gate." + "FlashInfer KDA prefill requires CUDA SM100 or SM103, bfloat16 " + "activations and recurrent state, head_dim=128, a bounded KDA " + "gate, and flashinfer-python 0.6.18 or newer." ) - if supported and backend != "triton": + if backend == "flashkda" and not flashkda_supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16 input, " + "bfloat16 or float32 recurrent state, head_dim=128, and a bounded " + "KDA gate." + ) + if flashinfer_supported and backend == "flashinfer": + logger.info_once("Using FlashInfer KDA prefill backend.") + return "flashinfer" + if flashkda_supported and backend in ("auto", "flashkda"): logger.info_once("Using FlashKDA KDA prefill backend.") return "flashkda" return "triton" @@ -369,7 +543,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") base_dtypes = MambaStateDtypeCalculator.kda_state_dtype( - self.model_config.dtype, self.cache_config.mamba_cache_dtype + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, ) if self.cache_config.use_kda_recoverssm: return MambaStateDtypeCalculator.append_kda_recoverssm_record( @@ -486,17 +662,25 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) # Keep a width-major copy for fused decode without changing the layout # consumed by the prefill and fallback decode kernels. - conv_state_dtype = self.get_state_dtype()[0] - decode_conv1d_weight = None - if is_fused_kda_decode_supported( + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2] + additional_config = vllm_config.additional_config + decode_backend = ( + additional_config.get("kda_decode_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + self.kda_decode_backend = resolve_kda_decode_backend( + decode_backend, self.local_num_heads, self.head_dim, self.conv_size, self.num_spec, vllm_config.model_config.dtype, conv_state_dtype, - ): - logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") + recurrent_state_dtype, + ) + decode_conv1d_weight = None + if self.kda_decode_backend != "triton": decode_conv1d_weight = torch.empty( 3, self.conv_size, @@ -533,7 +717,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): f"Got {self.gate_lower_bound}." ) - additional_config = vllm_config.additional_config backend = ( additional_config.get("kda_prefill_backend", "auto") if isinstance(additional_config, dict) @@ -543,11 +726,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): backend, self.head_dim, vllm_config.model_config.dtype, + recurrent_state_dtype, self.gate_lower_bound, ) self._flashkda_buffer_specs: ( tuple[tuple[tuple[int, ...], torch.dtype], ...] | None ) = None + self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = ( + None + ) if self.kda_prefill_backend == "flashkda": T = vllm_config.scheduler_config.max_num_batched_tokens N = vllm_config.scheduler_config.max_num_seqs @@ -561,6 +748,13 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): ((N, H, D, D), self.get_state_dtype()[1]), ((workspace_size,), torch.uint8), ) + elif self.kda_prefill_backend == "flashinfer": + T = vllm_config.scheduler_config.max_num_batched_tokens + H, D = self.local_num_heads, self.head_dim + self._flashinfer_kda_output_spec = ( + (1, T, H, D), + self.model_config.dtype, + ) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None @@ -725,30 +919,50 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): conv_state = conv_state.transpose(-1, -2) if ( - self.decode_conv1d_weight is not None + self.kda_decode_backend != "triton" + and self.decode_conv1d_weight is not None and self.decode_norm_weight is not None and not has_spec_decode and m.num_prefills == 0 and m.num_decodes > 0 ): assert non_spec_state_indices_tensor is not None - ops.fused_kda_decode( - x=mixed_qkv, - weight=self.decode_conv1d_weight, - bias=self.conv1d.bias, - conv_state=conv_state, - raw_g=g1, - raw_beta=beta, - A_log=self.A_log, - dt_bias=self.dt_bias, - state_indices=non_spec_state_indices_tensor[:num_actual_tokens], - state=recurrent_state, - out=core_attn_out[:, :num_actual_tokens], - lower_bound=self.gate_lower_bound, - output_gate=g2[:num_actual_tokens], - norm_weight=self.decode_norm_weight, - norm_eps=self.o_norm.eps, - ) + state_indices = non_spec_state_indices_tensor[:num_actual_tokens] + if self.kda_decode_backend == "flashinfer": + flashinfer_fused_kda_decode( + x=mixed_qkv, + weight=self.decode_conv1d_weight, + conv_state=conv_state, + raw_gate=g1, + raw_beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=state_indices, + state=recurrent_state, + output_gate=g2[:num_actual_tokens], + norm_weight=self.decode_norm_weight, + lower_bound=self.gate_lower_bound, + norm_eps=self.o_norm.eps, + output=core_attn_out[:, :num_actual_tokens], + ) + else: + ops.fused_kda_decode( + x=mixed_qkv, + weight=self.decode_conv1d_weight, + bias=self.conv1d.bias, + conv_state=conv_state, + raw_g=g1, + raw_beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=state_indices, + state=recurrent_state, + out=core_attn_out[:, :num_actual_tokens], + lower_bound=self.gate_lower_bound, + output_gate=g2[:num_actual_tokens], + norm_weight=self.decode_norm_weight, + norm_eps=self.o_norm.eps, + ) return conv_weights = self.conv1d.weight.view( @@ -991,6 +1205,35 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): final_state=final_state[: initial_state.shape[0]], workspace=workspace, ) + elif self.kda_prefill_backend == "flashinfer": + assert self.gate_lower_bound is not None + assert m.flashinfer_prefill_query_start_loc is not None + if q_ns.shape[1] > initial_state.shape[0]: + assert m.flashinfer_prefill_seq_order is not None + flashinfer_out = core_attn_out[:, : q_ns.shape[1]] + if has_spec_decode: + assert self._flashinfer_kda_output_spec is not None + (workspace_out,) = current_workspace_manager().get_simultaneous( + self._flashinfer_kda_output_spec + ) + flashinfer_out = workspace_out[:, : q_ns.shape[1]] + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashinfer_kda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=m.flashinfer_prefill_query_start_loc, + out=flashinfer_out, + seq_order=m.flashinfer_prefill_seq_order, + ) else: ( core_attn_out_non_spec, @@ -1009,7 +1252,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): use_qk_l2norm_in_kernel=True, cu_seqlens=non_spec_query_start_loc, ) - recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + recurrent_state[non_spec_state_indices_tensor] = ( + last_recurrent_state.to(recurrent_state.dtype) + ) else: # Pure non-speculative decode. assert non_spec_state_indices_tensor is not None @@ -1046,7 +1291,10 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif core_attn_out_non_spec is not None: - if self.kda_prefill_backend != "flashkda" or m.num_prefills == 0: + if ( + self.kda_prefill_backend not in ("flashkda", "flashinfer") + or m.num_prefills == 0 + ): # TODO: decode kernels write directly to core_attn_out core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ 0, :num_actual_tokens diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 1035f45948..21fdbea0b1 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -270,6 +270,8 @@ class KDACheckpointMetadata: @dataclass class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + flashinfer_prefill_query_start_loc: torch.Tensor | None = None + flashinfer_prefill_seq_order: torch.Tensor | None = None recoverssm_commit: KDARecoverSSMCommitMetadata | None = None recoverssm_context: "KDARecoverSSMCommitContext | None" = field( default=None, repr=False, compare=False @@ -318,6 +320,11 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + additional_config = vllm_config.additional_config + self.use_flashinfer_prefill = ( + isinstance(additional_config, dict) + and additional_config.get("kda_prefill_backend") == "flashinfer" + ) self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 self.recoverssm_num_accepted_tokens: torch.Tensor | None = None @@ -721,6 +728,20 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): align=align, ) + flashinfer_prefill_query_start_loc = None + flashinfer_prefill_seq_order = None + if self.use_flashinfer_prefill and num_prefills > 0: + assert non_spec_query_start_loc is not None + flashinfer_prefill_query_start_loc = non_spec_query_start_loc.to( + torch.int64 + ) + num_non_spec_requests = non_spec_query_start_loc.shape[0] - 1 + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + if num_non_spec_tokens > num_non_spec_requests: + flashinfer_prefill_seq_order = torch.argsort( + flashinfer_prefill_query_start_loc.diff(), descending=True + ).to(torch.int32) + return KimiK3KDAMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -738,6 +759,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, + flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc, + flashinfer_prefill_seq_order=flashinfer_prefill_seq_order, recoverssm_commit=recoverssm_commit, recoverssm_context=( self._get_recoverssm_context() diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 5c18348cf2..8bd04a12e6 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1649,7 +1649,9 @@ class KimiLinearForCausalLM( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, ...]: dtypes = MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, ) if vllm_config.cache_config.use_kda_recoverssm: dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 294e91391b..8c158dedb7 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -264,6 +264,14 @@ flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper( "flashinfer.decode", "xqa_batch_decode_with_kv_cache", ) +flashinfer_recurrent_kda = _lazy_import_wrapper( + "flashinfer.kda", + "recurrent_kda", +) +flashinfer_fused_kda_decode = _lazy_import_wrapper( + "flashinfer.kda_decode", + "fused_kda_decode", +) # Special case for autotune since it returns a context manager @@ -399,6 +407,28 @@ def has_flashinfer_bf16_fp4() -> bool: ) +@functools.cache +def has_flashinfer_recurrent_kda() -> bool: + """Return whether FlashInfer recurrent KDA prefill is available.""" + if not has_flashinfer(): + return False + mod = _get_submodule("flashinfer.kda") + return mod is not None and callable(getattr(mod, "recurrent_kda", None)) + + +@functools.cache +def has_flashinfer_fused_kda_decode() -> bool: + """Return whether FlashInfer fused KDA decode is available.""" + if not has_flashinfer(): + return False + mod = _get_submodule("flashinfer.kda_decode") + return ( + mod is not None + and bool(getattr(mod, "_FUSED_KDA_DECODE_AVAILABLE", False)) + and callable(getattr(mod, "fused_kda_decode", None)) + ) + + @functools.cache def has_flashinfer_trtllm_fused_moe() -> bool: """Return `True` if FlashInfer TRTLLM fused MoE is available.""" @@ -1261,6 +1291,8 @@ __all__ = [ "flashinfer_trtllm_batch_decode_with_kv_cache_mla", "flashinfer_trtllm_batch_decode_sparse_mla_dsv4", "flashinfer_xqa_batch_decode_with_kv_cache", + "flashinfer_recurrent_kda", + "flashinfer_fused_kda_decode", "autotune", "has_flashinfer_moe", "has_flashinfer_comm", @@ -1268,6 +1300,8 @@ __all__ = [ "has_flashinfer_nvlink_one_sided", "has_flashinfer_cutlass_fused_moe", "has_flashinfer_cutedsl_grouped_gemm_nt_masked", + "has_flashinfer_recurrent_kda", + "has_flashinfer_fused_kda_decode", "has_flashinfer_cutedsl_moe_nvfp4", "has_flashinfer_bf16_fp4", "has_flashinfer_b12x_moe",