vllm.models.kimi_k3.amd.ops.third_party.kda

Modules:

• chunk –
• chunk_intra –
• chunk_intra_token_parallel –
• fused_recurrent –

Functions:

• chunk_kda_with_fused_gate –

  Run chunk KDA from raw gate and beta projections.

• fused_kda_gate –

  Forward pass for KDA gate:

• fused_recurrent_kda –

  Run recurrent KDA from raw gate and beta inputs.

• fused_recurrent_kda_fwd –

  Launch recurrent KDA with dense inner dimensions and row strides.

• fused_recurrent_kda_packed_decode –

  Run one-token KDA decode directly from packed post-conv QKV.

chunk_kda_with_fused_gate(q, k, v, raw_g, raw_beta, A_log, g_bias, scale=None, initial_state=None, output_final_state=False, lower_bound=None, use_qk_l2norm_in_kernel=False, cu_seqlens=None, chunk_indices=None, chunk_offsets=None, **kwargs)

Run chunk KDA from raw gate and beta projections.

Source code in vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

def chunk_kda_with_fused_gate(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    raw_g: torch.Tensor,
    raw_beta: torch.Tensor,
    A_log: torch.Tensor,
    g_bias: torch.Tensor | None,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    output_final_state: bool = False,
    lower_bound: float | None = None,
    use_qk_l2norm_in_kernel: bool = False,
    cu_seqlens: torch.Tensor | None = None,
    chunk_indices: torch.Tensor | None = None,
    chunk_offsets: torch.Tensor | None = None,
    **kwargs,
):
    """Run chunk KDA from raw gate and beta projections."""
    if scale is None:
        scale = k.shape[-1] ** -0.5

    if use_qk_l2norm_in_kernel:
        q = l2norm_fwd(q.contiguous())
        k = l2norm_fwd(k.contiguous())

    o, final_state = chunk_kda_with_fused_gate_fwd(
        q=q,
        k=k,
        v=v.contiguous(),
        raw_g=raw_g.contiguous(),
        raw_beta=raw_beta,
        A_log=A_log,
        g_bias=g_bias,
        scale=scale,
        initial_state=initial_state.contiguous() if initial_state is not None else None,
        output_final_state=output_final_state,
        lower_bound=lower_bound,
        cu_seqlens=cu_seqlens,
        chunk_indices=chunk_indices,
        chunk_offsets=chunk_offsets,
    )
    return o, final_state

fused_kda_gate(g, A, head_k_dim, g_bias=None, beta=1.0, threshold=20.0, lower_bound=None)

Forward pass for KDA gate: input g: [..., H*D] param A: [H] or [1, 1, H, 1] beta: softplus beta parameter threshold: softplus threshold parameter return : [..., H, D]

Source code in vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

def fused_kda_gate(
    g: torch.Tensor,
    A: torch.Tensor,
    head_k_dim: int,
    g_bias: torch.Tensor | None = None,
    beta: float = 1.0,
    threshold: float = 20.0,
    lower_bound: float | None = None,
) -> torch.Tensor:
    """Forward pass for KDA gate:
    input g: [..., H*D]
    param A: [H] or [1, 1, H, 1]
    beta: softplus beta parameter
    threshold: softplus threshold parameter
    return  : [..., H, D]
    """
    orig_shape = g.shape[:-1]

    g = g.view(-1, g.shape[-1])
    T = g.shape[0]
    HD = g.shape[1]
    H = A.numel()
    assert H * head_k_dim == HD

    y = torch.empty_like(g, dtype=torch.float32)

    def grid(meta):
        return (cdiv(T, meta["BT"]), H)

    kda_gate_fwd_kernel[grid](
        g,
        A,
        y,
        g_bias,
        lower_bound or 0.0,
        beta,
        threshold,
        T,
        H,
        head_k_dim,
        BD=next_power_of_2(head_k_dim),
        HAS_BIAS=g_bias is not None,
        USE_LOWER_BOUND=lower_bound is not None,
    )

    y = y.view(*orig_shape, H, head_k_dim)
    return y

fused_recurrent_kda(q, k, v, raw_g, raw_beta, A_log, dt_bias, lower_bound, initial_state, cu_seqlens, ssm_state_indices, num_accepted_tokens=None, uniform_sequence_length=None, out=None, fuse_gate=None)

Run recurrent KDA from raw gate and beta inputs.

This vLLM wrapper applies the gate activation and beta sigmoid, selecting whether to materialize them before launching the recurrent kernel.

Source code in vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py

def fused_recurrent_kda(
    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 | None,
    lower_bound: float | None,
    initial_state: torch.Tensor,
    cu_seqlens: torch.Tensor,
    ssm_state_indices: torch.Tensor,
    num_accepted_tokens: torch.Tensor | None = None,
    uniform_sequence_length: int | None = None,
    out: torch.Tensor | None = None,
    fuse_gate: bool | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run recurrent KDA from raw gate and beta inputs.

    This vLLM wrapper applies the gate activation and beta sigmoid, selecting
    whether to materialize them before launching the recurrent kernel.
    """
    if fuse_gate is None:
        # gfx950: always fuse the gate.
        fuse_gate = True

    if fuse_gate:
        gate = raw_g
        beta = raw_beta
    else:
        gate, beta = _fused_kda_gate_beta(
            raw_g,
            raw_beta,
            A_log,
            dt_bias,
            lower_bound,
        )
    return fused_recurrent_kda_fwd(
        q=q,
        k=k,
        v=v,
        g=gate,
        beta=beta,
        scale=q.shape[-1] ** -0.5,
        initial_state=initial_state,
        inplace_final_state=True,
        cu_seqlens=cu_seqlens,
        ssm_state_indices=ssm_state_indices,
        num_accepted_tokens=num_accepted_tokens,
        uniform_sequence_length=uniform_sequence_length,
        use_qk_l2norm_in_kernel=True,
        A_log=A_log if fuse_gate else None,
        dt_bias=dt_bias if fuse_gate else None,
        lower_bound=lower_bound if fuse_gate else None,
        use_gate_in_kernel=fuse_gate,
        use_beta_sigmoid_in_kernel=fuse_gate,
        out=out,
    )

fused_recurrent_kda_fwd(q, k, v, g, beta, scale=None, initial_state=None, inplace_final_state=True, cu_seqlens=None, ssm_state_indices=None, num_accepted_tokens=None, uniform_sequence_length=None, use_qk_l2norm_in_kernel=True, A_log=None, dt_bias=None, lower_bound=None, use_gate_in_kernel=False, use_beta_sigmoid_in_kernel=False, out=None)

Launch recurrent KDA with dense inner dimensions and row strides.

Source code in vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py

def fused_recurrent_kda_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    inplace_final_state: bool = True,
    cu_seqlens: torch.Tensor | None = None,
    ssm_state_indices: torch.Tensor | None = None,
    num_accepted_tokens: torch.Tensor | None = None,
    uniform_sequence_length: int | None = None,
    use_qk_l2norm_in_kernel: bool = True,
    A_log: torch.Tensor | None = None,
    dt_bias: torch.Tensor | None = None,
    lower_bound: float | None = None,
    use_gate_in_kernel: bool = False,
    use_beta_sigmoid_in_kernel: bool = False,
    out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Launch recurrent KDA with dense inner dimensions and row strides."""
    B, T, H, K = q.shape
    V = v.shape[-1]
    assert B == 1 and k.shape == q.shape
    assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K)
    assert beta.shape == (B, T, H)
    assert initial_state is not None
    assert cu_seqlens is not None
    assert ssm_state_indices is not None
    assert inplace_final_state
    if out is None:
        out = torch.empty_like(v)
    assert out.shape == v.shape
    assert initial_state.shape[1:] == (H, V, K)
    assert ssm_state_indices.ndim in (1, 2)

    assert q.stride()[2:] == k.stride()[2:] == (K, 1)
    assert v.stride()[2:] == out.stride()[2:] == (V, 1)
    assert g.stride()[2:] == (K, 1)
    assert beta.stride(2) == 1
    assert q.stride(1) == k.stride(1) == v.stride(1)
    assert initial_state.stride()[1:] == (V * K, K, 1)
    N = cu_seqlens.numel() - 1
    if ssm_state_indices.ndim == 1:
        assert T == N
        assert num_accepted_tokens is None
    else:
        assert ssm_state_indices.stride(1) == 1
    if uniform_sequence_length is not None:
        assert num_accepted_tokens is not None
        assert uniform_sequence_length > 0
        assert N * uniform_sequence_length == T
    assert cu_seqlens.is_contiguous()
    if use_gate_in_kernel:
        assert A_log is not None and A_log.is_contiguous()
        assert dt_bias is None or dt_bias.is_contiguous()

    if scale is None:
        scale = K**-0.5

    is_spec_decoding = num_accepted_tokens is not None
    is_single_sequence = is_spec_decoding and N == 1
    head_sequences = H * N
    static_sequence_length = _get_static_sequence_length(
        total_tokens=T,
        head_sequences=head_sequences,
        is_single_sequence=is_single_sequence,
        uniform_sequence_length=uniform_sequence_length,
    )
    BV, num_warps, num_stages = _select_kda_launch_config(
        use_gate_in_kernel=use_gate_in_kernel,
        is_spec_decoding=is_spec_decoding,
        head_sequences=head_sequences,
        static_sequence_length=static_sequence_length,
        has_multi_token_sequence=T > N,
    )
    grid = (cdiv(V, BV) * N * H,)
    fused_recurrent_kda_fwd_kernel[grid](
        q=q,
        k=k,
        v=v,
        g=g,
        beta=beta,
        A_log=A_log,
        dt_bias=dt_bias,
        out=out,
        state=initial_state,
        cu_seqlens=cu_seqlens,
        state_indices=ssm_state_indices,
        num_accepted_tokens=num_accepted_tokens,
        lower_bound=lower_bound,
        scale=scale,
        N=N,
        T=T,
        H=H,
        K=K,
        V=V,
        BK=next_power_of_2(K),
        BV=BV,
        stride_qkv_token=q.stride(1),
        stride_g_token=g.stride(1),
        stride_beta_token=beta.stride(1),
        stride_out_token=out.stride(1),
        stride_state_token=initial_state.stride(0),
        stride_indices_seq=ssm_state_indices.stride(0),
        IS_SPEC_DECODING=num_accepted_tokens is not None,
        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
        USE_GATE_IN_KERNEL=use_gate_in_kernel,
        APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
        IS_SINGLE_SEQUENCE=is_single_sequence,
        SEQUENCE_LENGTH=static_sequence_length,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    return out, initial_state

fused_recurrent_kda_packed_decode(mixed_qkv, raw_g, raw_beta, A_log, dt_bias, lower_bound, initial_state, state_indices, scale=None)

Run one-token KDA decode directly from packed post-conv QKV.

Source code in vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py

def fused_recurrent_kda_packed_decode(
    mixed_qkv: torch.Tensor,
    raw_g: torch.Tensor,
    raw_beta: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    lower_bound: float | None,
    initial_state: torch.Tensor,
    state_indices: torch.Tensor,
    scale: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run one-token KDA decode directly from packed post-conv QKV."""
    if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
        raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.")
    if raw_g.ndim != 4 or raw_g.shape[0] != 1:
        raise ValueError("`raw_g` must have shape [1, B, H, K].")
    if raw_beta.ndim != 3 or raw_beta.shape[0] != 1:
        raise ValueError("`raw_beta` must have shape [1, B, H].")
    if initial_state.ndim != 4:
        raise ValueError("`initial_state` must have shape [cache, H, V, K].")
    _, H, V, K = initial_state.shape
    if raw_g.stride()[2:] != (K, 1):
        raise ValueError("`raw_g` must be contiguous within each token.")
    if raw_beta.stride(2) != 1:
        raise ValueError("`raw_beta` heads must be contiguous.")
    if initial_state.stride()[1:] != (V * K, K, 1):
        raise ValueError("`initial_state` must be contiguous within each cache slot.")
    if state_indices.ndim != 1:
        raise ValueError("`state_indices` must be one-dimensional.")
    if A_log.ndim != 1 or not A_log.is_contiguous():
        raise ValueError("`A_log` must be contiguous and one-dimensional.")
    if not dt_bias.is_contiguous():
        raise ValueError("`dt_bias` must be contiguous.")

    device = mixed_qkv.device
    if any(
        x.device != device
        for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices)
    ):
        raise ValueError("All packed KDA inputs must be on the same device.")

    B = mixed_qkv.shape[0]
    if raw_g.shape != (1, B, H, K):
        raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.")
    if raw_beta.shape != (1, B, H):
        raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.")
    if mixed_qkv.shape[1] != 2 * H * K + H * V:
        raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.")
    if A_log.numel() != H or dt_bias.numel() != H * K:
        raise ValueError("`A_log` or `dt_bias` has an incompatible shape.")
    if state_indices.shape[0] != B:
        raise ValueError("`state_indices` must contain one entry per token.")

    BK = next_power_of_2(K)
    BV = min(next_power_of_2(V), 32)
    if scale is None:
        scale = K**-0.5

    out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
    grid = (cdiv(V, BV), B * H)
    fused_recurrent_kda_packed_decode_kernel[grid](
        mixed_qkv=mixed_qkv,
        raw_g=raw_g,
        raw_beta=raw_beta,
        A_log=A_log,
        dt_bias=dt_bias,
        out=out,
        state=initial_state,
        state_indices=state_indices,
        lower_bound=lower_bound or 0.0,
        scale=scale,
        stride_mixed_token=mixed_qkv.stride(0),
        stride_g_token=raw_g.stride(1),
        stride_beta_token=raw_beta.stride(1),
        stride_state_token=initial_state.stride(0),
        stride_state_indices=state_indices.stride(0),
        H=H,
        K=K,
        V=V,
        BK=BK,
        BV=BV,
        SOFTPLUS_THRESHOLD=20.0,
        USE_LOWER_BOUND=lower_bound is not None,
        num_warps=4,
        num_stages=2,
    )
    return out, initial_state