| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm entry points for the fused Kimi-K3 KDA decode kernel. |
| 4 | |
| 5 | The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu`` |
| 6 | replaces, for a pure non-speculative decode batch, the three Triton launches |
| 7 | and two copies the AMD KDA layer otherwise runs per layer: the packed causal |
| 8 | conv1d update, the recurrent delta-rule step, and the gated output RMSNorm. |
| 9 | |
| 10 | The kernel wants a width-major conv weight and an fp32 norm weight, so both are |
| 11 | staged once at load time by the weight loaders below. |
| 12 | """ |
| 13 | |
| 14 | from collections.abc import Callable |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first |
| 20 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | # Head counts the kernel is instantiated for (Kimi-K3 has 96 KDA heads, so this |
| 25 | # covers TP1/2/4/8). |
| 26 | SUPPORTED_NUM_HEADS = (12, 24, 48, 96) |
| 27 | |
| 28 | |
| 29 | def is_fused_kda_decode_supported( |
| 30 | num_heads: int, |
| 31 | head_dim: int, |
| 32 | conv_width: int, |
| 33 | num_spec: int, |
| 34 | input_dtype: torch.dtype, |
| 35 | conv_state_dtype: torch.dtype, |
| 36 | ) -> bool: |
| 37 | """Whether the fused decode kernel can serve this layer on this device.""" |
| 38 | from vllm.platforms.rocm import on_gfx942, on_gfx950 |
| 39 | |
| 40 | if ( |
| 41 | num_heads not in SUPPORTED_NUM_HEADS |
| 42 | or head_dim != 128 |
| 43 | or conv_width != 4 |
| 44 | or num_spec != 0 |
| 45 | or input_dtype != torch.bfloat16 |
| 46 | or conv_state_dtype != torch.bfloat16 |
| 47 | or is_conv_state_dim_first() |
| 48 | or not hasattr(torch.ops._C, "fused_kda_decode") |
| 49 | ): |
| 50 | return False |
| 51 | # gfx950 (MI355X) and gfx942 (MI325X): both CDNA, sharing the wave64 / DPP / |
| 52 | # bf16 primitives the kernel relies on. |
| 53 | return on_gfx950() or on_gfx942() |
| 54 | |
| 55 | |
| 56 | def make_decode_conv1d_weight_loader( |
| 57 | dims: list[int], |
| 58 | tp_size: int, |
| 59 | tp_rank: int, |
| 60 | decode_conv1d_weight: torch.Tensor | None, |
| 61 | ) -> Callable[..., None]: |
| 62 | """Load the packed conv1d weight, mirroring a width-major fp32 copy. |
| 63 | |
| 64 | The fused kernel indexes the weight as ``[qkv, width, channel]`` so the |
| 65 | channel dimension is the contiguous one; the Triton prefill and fallback |
| 66 | decode kernels keep the ``[channel, width]`` layout. |
| 67 | """ |
| 68 | sharded_dims = [dim // tp_size for dim in dims] |
| 69 | |
| 70 | def weight_loader( |
| 71 | param: torch.Tensor, |
| 72 | loaded_weight: torch.Tensor, |
| 73 | loaded_shard_id: int, |
| 74 | ) -> None: |
| 75 | if loaded_weight.dim() == 2: |
| 76 | loaded_weight = loaded_weight.unsqueeze(1) |
| 77 | shard_size = sharded_dims[loaded_shard_id] |
| 78 | source_start = tp_rank * shard_size |
| 79 | target_start = sum(sharded_dims[:loaded_shard_id]) |
| 80 | loaded_shard = loaded_weight[source_start : source_start + shard_size] |
| 81 | param.data[target_start : target_start + shard_size].copy_(loaded_shard) |
| 82 | if decode_conv1d_weight is not None and not param.is_meta: |
| 83 | decode_conv1d_weight[loaded_shard_id].copy_( |
| 84 | loaded_shard.squeeze(1).transpose(0, 1) |
| 85 | ) |
| 86 | |
| 87 | return weight_loader |
| 88 | |
| 89 | |
| 90 | def make_decode_norm_weight_loader( |
| 91 | decode_norm_weight: torch.Tensor, |
| 92 | ) -> Callable[..., None]: |
| 93 | """Load the gated-norm weight, mirroring an fp32 copy for the kernel.""" |
| 94 | |
| 95 | def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 96 | default_weight_loader(param, loaded_weight) |
| 97 | if not param.is_meta: |
| 98 | decode_norm_weight.copy_(param.data) |
| 99 | |
| 100 | return weight_loader |