Skip to content

vllm.models.kimi_k3.amd.ops.kda_checkpoint

Conv half of the ROCm Kimi-K3 prefill checkpoint.

Classes:

KimiK3ROCmKDAPrefillCheckpointExporter dataclass

Bases: FlashKDAPrefillCheckpointExporter

Conv window only; fused_kda_chunk writes the recurrent half itself.

state_len is conv_kernel_size - 1, not the spec-widened conv row.

Source code in vllm/models/kimi_k3/amd/ops/kda_checkpoint.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
@dataclass(frozen=True)
class KimiK3ROCmKDAPrefillCheckpointExporter(FlashKDAPrefillCheckpointExporter):
    """Conv window only; `fused_kda_chunk` writes the recurrent half itself.

    `state_len` is `conv_kernel_size - 1`, not the spec-widened conv row.
    """

    def export(  # type: ignore[override]
        self,
        checkpoint: MambaPrefillCheckpointMetadata,
        *,
        raw_qkv: torch.Tensor,
        conv_state: torch.Tensor,
        cu_seqlens: torch.Tensor,
    ) -> None:
        assert self.state_len is not None
        # An empty recurrent checkpoint masks the kernel's recurrent half off.
        super().export(
            checkpoint,
            raw_qkv=raw_qkv,
            conv_state=conv_state,
            recurrent_checkpoint=conv_state[:, :0],
            recurrent_state=conv_state,
            cu_seqlens=cu_seqlens,
        )