Jiangyun ZhuandClaude Opus 5.5
badfcdbf91
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
15 lines
466 B
Python
15 lines
466 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from vllm.model_executor.layers.mamba.checkpoint import (
|
|
MambaPrefillCheckpointExporter,
|
|
)
|
|
|
|
|
|
def kda_prefill_checkpoint_alignment(backend: str) -> int | None:
|
|
return 16 if backend == "flashkda" else None
|
|
|
|
|
|
class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter):
|
|
"""Store FlashKDA recurrent and convolution checkpoint states."""
|