Skip to content

vllm.models.kimi_k3.amd.ops.kda_prefill

KDA prefill backend selection for ROCm.

The Kimi-K3 KDA layer calls :func:chunk_kda_prefill, which either runs the fused HIP kernels in kda_chunk or falls back to the vendored Triton chunk path.

Functions:

chunk_kda_prefill(q, k, v, raw_g, raw_beta, A_log, g_bias=None, 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, use_fused_chunk=False, out=None, checkpoint_state=None, checkpoint_offsets=None, checkpoint_state_indices=None, state_cache=None, state_indices=None, has_initial_state=None)

Run chunk KDA from raw gate and beta projections.

Parameters:

  • q

    (Tensor) –

    query tensor, [1, T, H, K].

  • k

    (Tensor) –

    key tensor, [1, T, H, K].

  • v

    (Tensor) –

    value tensor, [1, T, H, V].

  • raw_g

    (Tensor) –

    raw gate projection, before the activation.

  • raw_beta

    (Tensor) –

    raw beta projection, before the activation.

  • A_log

    (Tensor) –

    log of the per-head gate decay.

  • g_bias

    (Tensor | None, default: None ) –

    optional per-head gate bias.

  • scale

    (float | None, default: None ) –

    scale applied to the query-key products. Defaults to k.shape[-1] ** -0.5.

  • initial_state

    (Tensor | None, default: None ) –

    fp32 per-sequence initial recurrent state, or None.

  • output_final_state

    (bool, default: False ) –

    whether to return the final recurrent state.

  • lower_bound

    (float | None, default: None ) –

    optional floor applied to the gate.

  • use_qk_l2norm_in_kernel

    (bool, default: False ) –

    fold the q/k L2 norm into the kernel.

  • cu_seqlens

    (Tensor | None, default: None ) –

    int32 cumulative sequence lengths.

  • chunk_indices

    (Tensor | None, default: None ) –

    precomputed chunk indices for cu_seqlens.

  • chunk_offsets

    (Tensor | None, default: None ) –

    int32 per-sequence first chunk index.

  • use_fused_chunk

    (bool, default: False ) –

    request the two-kernel ROCm path. It is used only when every one of its preconditions holds; otherwise the Triton path runs unchanged.

  • out

    (Tensor | None, default: None ) –

    buffer the result must land in. Honoured by both backends, so the caller can hand in a slice of its own output and skip a copy.

  • checkpoint_state

    (Tensor | None, default: None ) –

    destination for mid-prefill recurrent state snapshots, letting a later prefix-cache hit resume from a mamba block boundary. See :func:fused_kda_chunk.

  • checkpoint_offsets

    (Tensor | None, default: None ) –

    per-sequence token offset to snapshot at, 0 for none.

  • checkpoint_state_indices

    (Tensor | None, default: None ) –

    optional per-sequence destination row.

  • state_cache

    (Tensor | None, default: None ) –

    the paged recurrent state. When given, the fused backend reads and writes it in place and neither a gather nor a scatter is needed around this call; the returned final state is None.

  • state_indices

    (Tensor | None, default: None ) –

    per-sequence cache row.

  • has_initial_state

    (Tensor | None, default: None ) –

    per-sequence flag; false starts from a zero state.

Returns:

  • tuple[Tensor, Tensor | None] –

    The output and, when requested, the final recurrent state.

Source code in vllm/models/kimi_k3/amd/ops/kda_prefill.py 
 27
 28
 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
 54
 55
 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
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def chunk_kda_prefill(
    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 = 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,
    use_fused_chunk: bool = False,
    out: torch.Tensor | None = None,
    checkpoint_state: torch.Tensor | None = None,
    checkpoint_offsets: torch.Tensor | None = None,
    checkpoint_state_indices: torch.Tensor | None = None,
    state_cache: torch.Tensor | None = None,
    state_indices: torch.Tensor | None = None,
    has_initial_state: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Run chunk KDA from raw gate and beta projections.

    Args:
        q: query tensor, ``[1, T, H, K]``.
        k: key tensor, ``[1, T, H, K]``.
        v: value tensor, ``[1, T, H, V]``.
        raw_g: raw gate projection, before the activation.
        raw_beta: raw beta projection, before the activation.
        A_log: log of the per-head gate decay.
        g_bias: optional per-head gate bias.
        scale: scale applied to the query-key products. Defaults to
            ``k.shape[-1] ** -0.5``.
        initial_state: fp32 per-sequence initial recurrent state, or ``None``.
        output_final_state: whether to return the final recurrent state.
        lower_bound: optional floor applied to the gate.
        use_qk_l2norm_in_kernel: fold the q/k L2 norm into the kernel.
        cu_seqlens: int32 cumulative sequence lengths.
        chunk_indices: precomputed chunk indices for ``cu_seqlens``.
        chunk_offsets: int32 per-sequence first chunk index.
        use_fused_chunk: request the two-kernel ROCm path. It is used only when
            every one of its preconditions holds; otherwise the Triton path
            runs unchanged.
        out: buffer the result must land in. Honoured by both backends, so the
            caller can hand in a slice of its own output and skip a copy.
        checkpoint_state: destination for mid-prefill recurrent state
            snapshots, letting a later prefix-cache hit resume from a mamba
            block boundary. See :func:`fused_kda_chunk`.
        checkpoint_offsets: per-sequence token offset to snapshot at, ``0``
            for none.
        checkpoint_state_indices: optional per-sequence destination row.
        state_cache: the paged recurrent state. When given, the fused backend
            reads and writes it in place and neither a gather nor a scatter is
            needed around this call; the returned final state is ``None``.
        state_indices: per-sequence cache row.
        has_initial_state: per-sequence flag; false starts from a zero state.

    Returns:
        The output and, when requested, the final recurrent state.

    """
    if scale is None:
        scale = k.shape[-1] ** -0.5

    # The fused prologue folds the q/k L2 norm and the gate activation in, so it
    # needs the raw projections and a bounded gate rather than the pre-normalized
    # tensors the Triton path takes.
    fused = (
        use_fused_chunk
        and use_qk_l2norm_in_kernel
        and cu_seqlens is not None
        and lower_bound is not None
        and g_bias is not None
        and can_use_fused_kda_chunk(k.shape[-1], v.shape[-1], k.dtype, FLA_CHUNK_SIZE)
    )

    if checkpoint_offsets is not None and not fused:
        raise NotImplementedError(
            "The KDA prefill checkpoint export needs kda_prefill_backend=fused for ROCm"
        )
    # Checked here rather than in the backend so both paths reject the same
    # arguments.
    if state_cache is not None:
        if initial_state is not None or output_final_state:
            raise ValueError("state_cache replaces initial_state/output_final_state")
        if state_indices is None or has_initial_state is None:
            raise ValueError("state_cache needs state_indices and has_initial_state")

    # Only the fused walk addresses the paged rows directly. The Triton path
    # gathers the rows it needs and scatters the results back here, so callers
    # pass the cache the same way for either backend.
    scatter_to: torch.Tensor | None = None
    if state_cache is not None and not fused:
        initial_state = gather_initial_states(
            state_cache, state_indices, has_initial_state
        )
        output_final_state = True
        scatter_to, state_cache = state_cache, None

    if fused:
        # Restated for the type checker; `fused` already implies all three.
        assert cu_seqlens is not None and lower_bound is not None
        assert g_bias is not None
        logger.info_once(
            "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel."
        )
        ws = fused_kda_prologue(
            q=q,
            k=k,
            v=v,
            raw_g=raw_g,
            raw_beta=raw_beta,
            A_log=A_log,
            dt_bias=g_bias,
            scale=scale,
            lower_bound=lower_bound,
            cu_seqlens=cu_seqlens,
            chunk_indices=chunk_indices,
        )
        return fused_kda_chunk(
            qg=ws["qg"],
            w=ws["w"],
            u=ws["u"],
            kg_t=ws["kg_t"],
            aqk=ws["aqk"],
            decay=ws["decay"],
            # v is dead by this point, so it doubles as the output buffer.
            out=out
            if out is not None
            else (
                v
                if v.is_contiguous()
                else torch.empty_like(v, memory_format=torch.contiguous_format)
            ),
            scale=scale,
            cu_seqlens=cu_seqlens,
            initial_state=initial_state,
            output_final_state=output_final_state,
            chunk_offsets=chunk_offsets,
            checkpoint_state=checkpoint_state,
            checkpoint_offsets=checkpoint_offsets,
            checkpoint_state_indices=checkpoint_state_indices,
            state_cache=state_cache,
            state_indices=state_indices,
            has_initial_state=has_initial_state,
        )

    o, final_state = chunk_kda_with_fused_gate(
        q=q,
        k=k,
        v=v,
        raw_g=raw_g,
        raw_beta=raw_beta,
        A_log=A_log,
        g_bias=g_bias,
        scale=scale,
        initial_state=initial_state,
        output_final_state=output_final_state,
        lower_bound=lower_bound,
        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
        cu_seqlens=cu_seqlens,
        chunk_indices=chunk_indices,
        chunk_offsets=chunk_offsets,
    )
    if out is not None and o.data_ptr() != out.data_ptr():
        out.copy_(o)
        o = out
    if scatter_to is not None:
        assert state_indices is not None
        scatter_to[state_indices.long()] = final_state
        return o, None
    return o, final_state