| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Warm up Kimi-K3 Triton kernels.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import sys |
| 8 | from typing import TYPE_CHECKING |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | if TYPE_CHECKING: |
| 16 | from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention |
| 17 | from vllm.v1.worker.gpu_worker import Worker |
| 18 | |
| 19 | logger = init_logger(__name__) |
| 20 | |
| 21 | |
| 22 | def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None: |
| 23 | # Kimi model construction already imports kda. Avoid importing it here. |
| 24 | kda = sys.modules.get("vllm.models.kimi_k3.nvidia.kda") |
| 25 | if kda is None: |
| 26 | return None |
| 27 | |
| 28 | compilation_config = getattr( |
| 29 | worker.model_runner, |
| 30 | "compilation_config", |
| 31 | None, |
| 32 | ) |
| 33 | static_context = getattr(compilation_config, "static_forward_context", None) |
| 34 | if not isinstance(static_context, dict): |
| 35 | return None |
| 36 | return next( |
| 37 | ( |
| 38 | layer |
| 39 | for layer in static_context.values() |
| 40 | if isinstance(layer, kda.KimiK3DeltaAttention) |
| 41 | ), |
| 42 | None, |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def _warm_attn_res(worker: Worker) -> None: |
| 47 | from vllm.models.kimi_k3.nvidia.ops.attn_res import ( |
| 48 | attn_res, |
| 49 | get_attn_res_triton_warmup_profiles, |
| 50 | ) |
| 51 | |
| 52 | config = worker.model_config.hf_text_config |
| 53 | block_size = getattr(config, "attn_res_block_size", None) |
| 54 | if block_size is None: |
| 55 | return |
| 56 | |
| 57 | hidden_size = int(config.hidden_size) |
| 58 | max_blocks = (int(config.num_hidden_layers) + block_size - 1) // block_size |
| 59 | if max_blocks < 2: |
| 60 | return |
| 61 | |
| 62 | dtype = worker.model_config.dtype |
| 63 | device = torch.device("cuda") |
| 64 | eps = float(config.rms_norm_eps) |
| 65 | prefix = torch.zeros((1, hidden_size), dtype=dtype, device=device) |
| 66 | delta = torch.zeros_like(prefix) |
| 67 | blocks = torch.zeros( |
| 68 | (1, max_blocks, hidden_size), |
| 69 | dtype=dtype, |
| 70 | device=device, |
| 71 | ) |
| 72 | norm_weight = torch.zeros(hidden_size, dtype=dtype, device=device) |
| 73 | qk_weight = torch.zeros_like(norm_weight) |
| 74 | output_norm_weight = torch.zeros_like(norm_weight) |
| 75 | |
| 76 | for ( |
| 77 | num_blocks, |
| 78 | has_delta, |
| 79 | block_write_idx, |
| 80 | apply_output_norm, |
| 81 | ) in get_attn_res_triton_warmup_profiles(max_blocks): |
| 82 | attn_res( |
| 83 | prefix, |
| 84 | delta if has_delta else None, |
| 85 | blocks, |
| 86 | norm_weight, |
| 87 | qk_weight, |
| 88 | output_norm_weight if apply_output_norm else None, |
| 89 | num_blocks=num_blocks, |
| 90 | block_write_idx=block_write_idx, |
| 91 | eps=eps, |
| 92 | output_norm_eps=eps if apply_output_norm else 0.0, |
| 93 | ) |
| 94 | |
| 95 | |
| 96 | def _warm_recurrent_kda( |
| 97 | layer: KimiK3DeltaAttention, |
| 98 | input_dtype: torch.dtype, |
| 99 | ) -> None: |
| 100 | from vllm.models.kimi_k3.nvidia.ops.third_party.kda.fused_recurrent import ( |
| 101 | fused_recurrent_kda, |
| 102 | get_fused_recurrent_kda_fwd_warmup_profiles, |
| 103 | ) |
| 104 | |
| 105 | num_speculative_tokens = int(layer.num_spec) |
| 106 | # fused_recurrent_kda_fwd_kernel is only used by speculative decode. |
| 107 | if num_speculative_tokens <= 0: |
| 108 | return |
| 109 | |
| 110 | kv_cache = layer.kv_cache |
| 111 | if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2: |
| 112 | return |
| 113 | state = kv_cache[1] |
| 114 | if not isinstance(state, torch.Tensor) or not state.numel(): |
| 115 | return |
| 116 | |
| 117 | logger.info("Warming up Kimi-K3 speculative KDA kernels.") |
| 118 | h = int(layer.local_num_heads) |
| 119 | d = int(layer.head_dim) |
| 120 | tokens_per_sequence = num_speculative_tokens + 1 |
| 121 | for num_sequences in get_fused_recurrent_kda_fwd_warmup_profiles(h): |
| 122 | num_tokens = num_sequences * tokens_per_sequence |
| 123 | packed_qkv = torch.empty( |
| 124 | (num_tokens, 3 * h * d), |
| 125 | dtype=input_dtype, |
| 126 | device=state.device, |
| 127 | ) |
| 128 | q, k, v = ( |
| 129 | tensor.view(1, num_tokens, h, d) |
| 130 | for tensor in packed_qkv.split(h * d, dim=-1) |
| 131 | ) |
| 132 | fused_recurrent_kda( |
| 133 | q=q, |
| 134 | k=k, |
| 135 | v=v, |
| 136 | raw_g=torch.empty( |
| 137 | (1, num_tokens, h, d), |
| 138 | dtype=input_dtype, |
| 139 | device=state.device, |
| 140 | ), |
| 141 | raw_beta=torch.empty( |
| 142 | (1, num_tokens, h), |
| 143 | dtype=input_dtype, |
| 144 | device=state.device, |
| 145 | ), |
| 146 | A_log=layer.A_log, |
| 147 | dt_bias=layer.dt_bias, |
| 148 | lower_bound=layer.gate_lower_bound, |
| 149 | initial_state=state[:1], |
| 150 | cu_seqlens=torch.arange( |
| 151 | 0, |
| 152 | num_tokens + 1, |
| 153 | tokens_per_sequence, |
| 154 | dtype=torch.int32, |
| 155 | device=state.device, |
| 156 | ), |
| 157 | ssm_state_indices=torch.zeros( |
| 158 | (num_sequences, tokens_per_sequence), |
| 159 | dtype=torch.int32, |
| 160 | device=state.device, |
| 161 | ), |
| 162 | num_accepted_tokens=torch.ones( |
| 163 | num_sequences, |
| 164 | dtype=torch.int32, |
| 165 | device=state.device, |
| 166 | ), |
| 167 | out=torch.empty( |
| 168 | (1, num_tokens, h, d), |
| 169 | dtype=input_dtype, |
| 170 | device=state.device, |
| 171 | ), |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | @torch.inference_mode() |
| 176 | def kimi_k3_triton_warmup(worker: Worker) -> None: |
| 177 | """Warm Kimi-K3 Triton kernels reachable by this server.""" |
| 178 | if not current_platform.is_cuda(): |
| 179 | return |
| 180 | |
| 181 | layer = _get_kda_layer(worker) |
| 182 | if layer is None: |
| 183 | return |
| 184 | |
| 185 | _warm_attn_res(worker) |
| 186 | _warm_recurrent_kda(layer, worker.model_config.dtype) |