File size: 1,872 Bytes
ff7679d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 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 42 43 44 45 46 47 48 49 | # 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
|