| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """KDA prefill backend selection for ROCm. |
| 4 | |
| 5 | The Kimi-K3 KDA layer calls :func:`chunk_kda_prefill`, which either runs the |
| 6 | fused HIP kernels in ``kda_chunk`` or falls back to the vendored Triton chunk |
| 7 | path. |
| 8 | """ |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( |
| 14 | gather_initial_states, |
| 15 | ) |
| 16 | from vllm.models.kimi_k3.amd.ops.kda_chunk import ( |
| 17 | can_use_fused_kda_chunk, |
| 18 | fused_kda_chunk, |
| 19 | fused_kda_prologue, |
| 20 | ) |
| 21 | from vllm.models.kimi_k3.amd.ops.third_party.kda import chunk_kda_with_fused_gate |
| 22 | from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | |
| 27 | def chunk_kda_prefill( |
| 28 | q: torch.Tensor, |
| 29 | k: torch.Tensor, |
| 30 | v: torch.Tensor, |
| 31 | raw_g: torch.Tensor, |
| 32 | raw_beta: torch.Tensor, |
| 33 | A_log: torch.Tensor, |
| 34 | g_bias: torch.Tensor | None = None, |
| 35 | scale: float | None = None, |
| 36 | initial_state: torch.Tensor | None = None, |
| 37 | output_final_state: bool = False, |
| 38 | lower_bound: float | None = None, |
| 39 | use_qk_l2norm_in_kernel: bool = False, |
| 40 | cu_seqlens: torch.Tensor | None = None, |
| 41 | chunk_indices: torch.Tensor | None = None, |
| 42 | chunk_offsets: torch.Tensor | None = None, |
| 43 | use_fused_chunk: bool = False, |
| 44 | out: torch.Tensor | None = None, |
| 45 | checkpoint_state: torch.Tensor | None = None, |
| 46 | checkpoint_offsets: torch.Tensor | None = None, |
| 47 | checkpoint_state_indices: torch.Tensor | None = None, |
| 48 | state_cache: torch.Tensor | None = None, |
| 49 | state_indices: torch.Tensor | None = None, |
| 50 | has_initial_state: torch.Tensor | None = None, |
| 51 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 52 | """Run chunk KDA from raw gate and beta projections. |
| 53 | |
| 54 | Args: |
| 55 | q: query tensor, ``[1, T, H, K]``. |
| 56 | k: key tensor, ``[1, T, H, K]``. |
| 57 | v: value tensor, ``[1, T, H, V]``. |
| 58 | raw_g: raw gate projection, before the activation. |
| 59 | raw_beta: raw beta projection, before the activation. |
| 60 | A_log: log of the per-head gate decay. |
| 61 | g_bias: optional per-head gate bias. |
| 62 | scale: scale applied to the query-key products. Defaults to |
| 63 | ``k.shape[-1] ** -0.5``. |
| 64 | initial_state: fp32 per-sequence initial recurrent state, or ``None``. |
| 65 | output_final_state: whether to return the final recurrent state. |
| 66 | lower_bound: optional floor applied to the gate. |
| 67 | use_qk_l2norm_in_kernel: fold the q/k L2 norm into the kernel. |
| 68 | cu_seqlens: int32 cumulative sequence lengths. |
| 69 | chunk_indices: precomputed chunk indices for ``cu_seqlens``. |
| 70 | chunk_offsets: int32 per-sequence first chunk index. |
| 71 | use_fused_chunk: request the two-kernel ROCm path. It is used only when |
| 72 | every one of its preconditions holds; otherwise the Triton path |
| 73 | runs unchanged. |
| 74 | out: buffer the result must land in. Honoured by both backends, so the |
| 75 | caller can hand in a slice of its own output and skip a copy. |
| 76 | checkpoint_state: destination for mid-prefill recurrent state |
| 77 | snapshots, letting a later prefix-cache hit resume from a mamba |
| 78 | block boundary. See :func:`fused_kda_chunk`. |
| 79 | checkpoint_offsets: per-sequence token offset to snapshot at, ``0`` |
| 80 | for none. |
| 81 | checkpoint_state_indices: optional per-sequence destination row. |
| 82 | state_cache: the paged recurrent state. When given, the fused backend |
| 83 | reads and writes it in place and neither a gather nor a scatter is |
| 84 | needed around this call; the returned final state is ``None``. |
| 85 | state_indices: per-sequence cache row. |
| 86 | has_initial_state: per-sequence flag; false starts from a zero state. |
| 87 | |
| 88 | Returns: |
| 89 | The output and, when requested, the final recurrent state. |
| 90 | |
| 91 | """ |
| 92 | if scale is None: |
| 93 | scale = k.shape[-1] ** -0.5 |
| 94 | |
| 95 | # The fused prologue folds the q/k L2 norm and the gate activation in, so it |
| 96 | # needs the raw projections and a bounded gate rather than the pre-normalized |
| 97 | # tensors the Triton path takes. |
| 98 | fused = ( |
| 99 | use_fused_chunk |
| 100 | and use_qk_l2norm_in_kernel |
| 101 | and cu_seqlens is not None |
| 102 | and lower_bound is not None |
| 103 | and g_bias is not None |
| 104 | and can_use_fused_kda_chunk(k.shape[-1], v.shape[-1], k.dtype, FLA_CHUNK_SIZE) |
| 105 | ) |
| 106 | |
| 107 | if checkpoint_offsets is not None and not fused: |
| 108 | raise NotImplementedError( |
| 109 | "The KDA prefill checkpoint export needs kda_prefill_backend=fused for ROCm" |
| 110 | ) |
| 111 | # Checked here rather than in the backend so both paths reject the same |
| 112 | # arguments. |
| 113 | if state_cache is not None: |
| 114 | if initial_state is not None or output_final_state: |
| 115 | raise ValueError("state_cache replaces initial_state/output_final_state") |
| 116 | if state_indices is None or has_initial_state is None: |
| 117 | raise ValueError("state_cache needs state_indices and has_initial_state") |
| 118 | |
| 119 | # Only the fused walk addresses the paged rows directly. The Triton path |
| 120 | # gathers the rows it needs and scatters the results back here, so callers |
| 121 | # pass the cache the same way for either backend. |
| 122 | scatter_to: torch.Tensor | None = None |
| 123 | if state_cache is not None and not fused: |
| 124 | initial_state = gather_initial_states( |
| 125 | state_cache, state_indices, has_initial_state |
| 126 | ) |
| 127 | output_final_state = True |
| 128 | scatter_to, state_cache = state_cache, None |
| 129 | |
| 130 | if fused: |
| 131 | # Restated for the type checker; `fused` already implies all three. |
| 132 | assert cu_seqlens is not None and lower_bound is not None |
| 133 | assert g_bias is not None |
| 134 | logger.info_once( |
| 135 | "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel." |
| 136 | ) |
| 137 | ws = fused_kda_prologue( |
| 138 | q=q, |
| 139 | k=k, |
| 140 | v=v, |
| 141 | raw_g=raw_g, |
| 142 | raw_beta=raw_beta, |
| 143 | A_log=A_log, |
| 144 | dt_bias=g_bias, |
| 145 | scale=scale, |
| 146 | lower_bound=lower_bound, |
| 147 | cu_seqlens=cu_seqlens, |
| 148 | chunk_indices=chunk_indices, |
| 149 | ) |
| 150 | return fused_kda_chunk( |
| 151 | qg=ws["qg"], |
| 152 | w=ws["w"], |
| 153 | u=ws["u"], |
| 154 | kg_t=ws["kg_t"], |
| 155 | aqk=ws["aqk"], |
| 156 | decay=ws["decay"], |
| 157 | # v is dead by this point, so it doubles as the output buffer. |
| 158 | out=out |
| 159 | if out is not None |
| 160 | else ( |
| 161 | v |
| 162 | if v.is_contiguous() |
| 163 | else torch.empty_like(v, memory_format=torch.contiguous_format) |
| 164 | ), |
| 165 | scale=scale, |
| 166 | cu_seqlens=cu_seqlens, |
| 167 | initial_state=initial_state, |
| 168 | output_final_state=output_final_state, |
| 169 | chunk_offsets=chunk_offsets, |
| 170 | checkpoint_state=checkpoint_state, |
| 171 | checkpoint_offsets=checkpoint_offsets, |
| 172 | checkpoint_state_indices=checkpoint_state_indices, |
| 173 | state_cache=state_cache, |
| 174 | state_indices=state_indices, |
| 175 | has_initial_state=has_initial_state, |
| 176 | ) |
| 177 | |
| 178 | o, final_state = chunk_kda_with_fused_gate( |
| 179 | q=q, |
| 180 | k=k, |
| 181 | v=v, |
| 182 | raw_g=raw_g, |
| 183 | raw_beta=raw_beta, |
| 184 | A_log=A_log, |
| 185 | g_bias=g_bias, |
| 186 | scale=scale, |
| 187 | initial_state=initial_state, |
| 188 | output_final_state=output_final_state, |
| 189 | lower_bound=lower_bound, |
| 190 | use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, |
| 191 | cu_seqlens=cu_seqlens, |
| 192 | chunk_indices=chunk_indices, |
| 193 | chunk_offsets=chunk_offsets, |
| 194 | ) |
| 195 | if out is not None and o.data_ptr() != out.data_ptr(): |
| 196 | out.copy_(o) |
| 197 | o = out |
| 198 | if scatter_to is not None: |
| 199 | assert state_indices is not None |
| 200 | scatter_to[state_indices.long()] = final_state |
| 201 | return o, None |
| 202 | return o, final_state |