Skip to content

vllm.models.kimi_k3.amd.ops.kda_decode

ROCm entry points for the fused Kimi-K3 KDA decode kernel.

The kernel in csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu replaces, for a pure non-speculative decode batch, the three Triton launches and two copies the AMD KDA layer otherwise runs per layer: the packed causal conv1d update, the recurrent delta-rule step, and the gated output RMSNorm.

The kernel wants a width-major conv weight and an fp32 norm weight, so both are staged once at load time by the weight loaders below.

Functions:

is_fused_kda_decode_supported(num_heads, head_dim, conv_width, num_spec, input_dtype, conv_state_dtype)

Whether the fused decode kernel can serve this layer on this device.

Source code in vllm/models/kimi_k3/amd/ops/kda_decode.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def is_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,
) -> bool:
    """Whether the fused decode kernel can serve this layer on this device."""
    from vllm.platforms.rocm import on_gfx942, on_gfx950

    if (
        num_heads not in SUPPORTED_NUM_HEADS
        or head_dim != 128
        or conv_width != 4
        or num_spec != 0
        or input_dtype != torch.bfloat16
        or conv_state_dtype != torch.bfloat16
        or is_conv_state_dim_first()
        or not hasattr(torch.ops._C, "fused_kda_decode")
    ):
        return False
    # gfx950 (MI355X) and gfx942 (MI325X): both CDNA, sharing the wave64 / DPP /
    # bf16 primitives the kernel relies on.
    return on_gfx950() or on_gfx942()

make_decode_conv1d_weight_loader(dims, tp_size, tp_rank, decode_conv1d_weight)

Load the packed conv1d weight, mirroring a width-major fp32 copy.

The fused kernel indexes the weight as [qkv, width, channel] so the channel dimension is the contiguous one; the Triton prefill and fallback decode kernels keep the [channel, width] layout.

Source code in vllm/models/kimi_k3/amd/ops/kda_decode.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def make_decode_conv1d_weight_loader(
    dims: list[int],
    tp_size: int,
    tp_rank: int,
    decode_conv1d_weight: torch.Tensor | None,
) -> Callable[..., None]:
    """Load the packed conv1d weight, mirroring a width-major fp32 copy.

    The fused kernel indexes the weight as ``[qkv, width, channel]`` so the
    channel dimension is the contiguous one; the Triton prefill and fallback
    decode kernels keep the ``[channel, width]`` layout.
    """
    sharded_dims = [dim // tp_size for dim in dims]

    def weight_loader(
        param: torch.Tensor,
        loaded_weight: torch.Tensor,
        loaded_shard_id: int,
    ) -> None:
        if loaded_weight.dim() == 2:
            loaded_weight = loaded_weight.unsqueeze(1)
        shard_size = sharded_dims[loaded_shard_id]
        source_start = tp_rank * shard_size
        target_start = sum(sharded_dims[:loaded_shard_id])
        loaded_shard = loaded_weight[source_start : source_start + shard_size]
        param.data[target_start : target_start + shard_size].copy_(loaded_shard)
        if decode_conv1d_weight is not None and not param.is_meta:
            decode_conv1d_weight[loaded_shard_id].copy_(
                loaded_shard.squeeze(1).transpose(0, 1)
            )

    return weight_loader

make_decode_norm_weight_loader(decode_norm_weight)

Load the gated-norm weight, mirroring an fp32 copy for the kernel.

Source code in vllm/models/kimi_k3/amd/ops/kda_decode.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def make_decode_norm_weight_loader(
    decode_norm_weight: torch.Tensor,
) -> Callable[..., None]:
    """Load the gated-norm weight, mirroring an fp32 copy for the kernel."""

    def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
        default_weight_loader(param, loaded_weight)
        if not param.is_meta:
            decode_norm_weight.copy_(param.data)

    return weight_loader