| |
| """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 |
|
|
