# SPDX-License-Identifier: Apache-2.0
"""Raon-local scalar KV loading for its native MLA attention implementation.

ModelOpt only recognizes its generic Attention/MLAAttention classes. Raon owns
an independent AttentionLayerBase implementation, so adapt only its explicitly
calibrated MXFP8 KV path; projection and expert quantization remain unchanged.
"""
import weakref

from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod


class _RaonMLAKVQuantConfig:
    def __init__(self, config, attention):
        self._config = config
        self._attention = weakref.ref(attention)

    def __getattr__(self, name):
        return getattr(self._config, name)

    def get_quant_method(self, layer, prefix):
        if layer is self._attention():
            return BaseKVCacheMethod(self._config)
        return self._config.get_quant_method(layer, prefix)


def mla_kv_quant_config(config, attention):
    """Preserve existing behavior unless the MXFP8 artifact declares FP8 KV."""
    if config is None or config.get_name() != "modelopt_mxfp8":
        return config
    algorithm = getattr(config, "kv_cache_quant_algo", None)
    if algorithm is None:
        return config
    if algorithm != "FP8":
        raise ValueError(f"Raon MLA requires FP8 KV calibration, got {algorithm!r}")
    return _RaonMLAKVQuantConfig(config, attention)


def dequantize_context_latent(gathered, scale, fp8_prefill, cache_dtype):
    """Restore a raw plain-FP8 gather; other gather paths already dequantize.

    Both the normalized latent and positional tail share the cache scale. Keep
    restoration in FP32 until the caller's projection/attention casts, avoiding
    an extra FP8 round trip before a BF16 projection.
    """
    if fp8_prefill and cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"):
        return gathered.float() * scale
    return gathered