################################################################################
# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
################################################################################

from typing import TYPE_CHECKING

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from fastcore.basics import patch_to

from vllm.logger import init_logger
from vllm.model_executor.models.config import DeepseekV32ForCausalLM

if TYPE_CHECKING:

    from vllm.config import VllmConfig

logger = init_logger(__name__)


@patch_to(DeepseekV32ForCausalLM)
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
    """
    Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32
    """
    hf_config = vllm_config.model_config.hf_config

    # Mirror the check in vllm/model_executor/models/deepseek_v2.py
    is_v32 = hasattr(hf_config, "index_topk")
    assert is_v32

    # For DeepSeekV3.2, a custom fp8 format is used when fp8 kv-cache is enabled.
    cache_config = vllm_config.cache_config
    if cache_config.cache_dtype.startswith("fp8"):
        cache_config.cache_dtype = "fp8_ds_mla"
        logger.info("Using custom fp8 kv-cache format for DeepSeekV3.2")
    if cache_config.cache_dtype == "bfloat16":
        cache_config.cache_dtype = "auto"
    logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")