| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Conv half of the ROCm Kimi-K3 prefill checkpoint.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.layers.mamba.checkpoint import ( |
| 10 | MambaPrefillCheckpointMetadata, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.mamba.kda_checkpoint import ( |
| 13 | FlashKDAPrefillCheckpointExporter, |
| 14 | ) |
| 15 | |
| 16 | |
| 17 | @dataclass(frozen=True) |
| 18 | class KimiK3ROCmKDAPrefillCheckpointExporter(FlashKDAPrefillCheckpointExporter): |
| 19 | """Conv window only; `fused_kda_chunk` writes the recurrent half itself. |
| 20 | |
| 21 | `state_len` is `conv_kernel_size - 1`, not the spec-widened conv row. |
| 22 | """ |
| 23 | |
| 24 | def export( # type: ignore[override] |
| 25 | self, |
| 26 | checkpoint: MambaPrefillCheckpointMetadata, |
| 27 | *, |
| 28 | raw_qkv: torch.Tensor, |
| 29 | conv_state: torch.Tensor, |
| 30 | cu_seqlens: torch.Tensor, |
| 31 | ) -> None: |
| 32 | assert self.state_len is not None |
| 33 | # An empty recurrent checkpoint masks the kernel's recurrent half off. |
| 34 | super().export( |
| 35 | checkpoint, |
| 36 | raw_qkv=raw_qkv, |
| 37 | conv_state=conv_state, |
| 38 | recurrent_checkpoint=conv_state[:, :0], |
| 39 | recurrent_state=conv_state, |
| 40 | cu_seqlens=cu_seqlens, |
| 41 | ) |